Update app.py

app.py

@@ -2,6 +2,8 @@ from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from sentence_transformers import SentenceTransformer, util
 from transformers import pipeline
+from transformers import T5Tokenizer, T5ForConditionalGeneration
+
 
 # Initialize FastAPI app
 app = FastAPI()
@@ -11,6 +13,9 @@ model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
 question_model = "deepset/tinyroberta-squad2"
 nlp = pipeline('question-answering', model=question_model, tokenizer=question_model)
 
+t5tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-large")
+t5model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-large")
+
 # Define request models
 class ModifyQueryRequest(BaseModel):
     query_string: str
@@ -20,6 +25,13 @@ class AnswerQuestionRequest(BaseModel):
     context: list
     locations: list
 
+class T5QuestionRequest(BaseModel):
+    question: str
+    context: list
+
+class T5Response(BaseModel):
+    answer: str
+
 # Define response models (if needed)
 class ModifyQueryResponse(BaseModel):
     embeddings: list
@@ -63,6 +75,15 @@ async def answer_question(request: AnswerQuestionRequest):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+@app.post("/t5answer", response_model=T5Response)
+async def t5answer(request: T5QuestionRequest):
+    input_text = request.question + ":" + request.context
+    input_ids = t5tokenizer(input_text, return_tensors="pt").input_ids
+    outputs = t5model.generate(input_ids)
+    resp = t5tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return T5Response(answer = resp)
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=8000)
+