| from fastapi import FastAPI |
| from pydantic import BaseModel |
| from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM |
| from huggingface_hub import login |
| import os |
|
|
| |
| login(token=os.getenv("HF_TOKEN")) |
|
|
| app = FastAPI() |
|
|
| pipe = pipeline("text-generation", model="mistralai/Mistral-Small-24B-Instruct-2501") |
| tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-Small-24B-Instruct-2501") |
| model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-Small-24B-Instruct-2501") |
|
|
| class ChatRequest(BaseModel): |
| messages: list |
|
|
| @app.get("/") |
| def read_root(): |
| return {"message": "Waredocs LLM API is running!"} |
|
|
| @app.get("/ask") |
| def ask_question(prompt: str): |
| messages = [{"role": "user", "content": prompt}] |
| result = pipe(prompt, max_length=200) |
| return result[0] |
