File size: 825 Bytes
2c3cdb2 8ec865f 2af189d ec34e4a 2c3cdb2 2af189d 2c3cdb2 8ec865f 2c3cdb2 c976b59 2af189d ec34e4a 2af189d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 | from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
from huggingface_hub import login
import os
# Authenticate with Hugging Face
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] |