Update main.py

main.py

@@ -1,24 +1,28 @@
 from fastapi import FastAPI
-from fastapi.responses import JSONResponse
+from models import sa_pipeline, query_index
 from pydantic import BaseModel
-from transformers import pipeline
-import uvicorn
 
-app = FastAPI()
+class QARequest(BaseModel):
+    question: str
 
+class QAResponse(BaseModel):
+    answer: str
 
-class Payload(BaseModel):
+class SentimentRequest(BaseModel):
     text: str
 
+class SentimentResponse(BaseModel):
+    label: str
+    score: float
 
-async def classify_text(text):
-    pipe = pipeline("text-classification",
-                    model="SamLowe/roberta-base-go_emotions")
-    return pipe(text)
+app = FastAPI()
 
+@app.post('/question-answering', response_model=QAResponse)
+def query(query: QARequest):
+    data = query.dict()
+    return {'answer': query_index(data['question'])}
 
-@app.post("/classify/")
-async def classify(payload: Payload):
-    result = await classify_text(payload.text)
-    return JSONResponse({"data": result})
-    
+@app.post('/sentiment-analysis', response_model=SentimentResponse)
+def query(query: SentimentRequest):
+    data = query.dict()
+    return sa_pipeline(data['text'])[0]