from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline
import uvicorn

# Load model (small QA model to fit free tier)
qa_pipeline = pipeline("question-answering", model="distilbert-base-uncased-distilled-squad")





# Define request format
class QARequest(BaseModel):
    question: str
    context: str

# Initialize FastAPI
app = FastAPI()

@app.get("/")
def home():
    return {"message": "Chatbot API is running 🚀"}

@app.post("/chat")
def chat(request: QARequest):
    result = qa_pipeline(question=request.question, context=request.context)
    return {"answer": result["answer"], "score": result["score"]}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)