LLM / main.py
hollywoodfrancis's picture
Create main.py
1112483 verified
raw
history blame contribute delete
351 Bytes
from fastapi import FastAPI
from transformers import pipeline
app = FastAPI()
# Load your model
model_name = "your-username/your-model-name"
qa_pipeline = pipeline("text-generation", model=model_name)
@app.post("/ask")
def ask(question: str):
response = qa_pipeline(question, max_length=100)
return {"answer": response[0]["generated_text"]}