Update app.py

app.py

@@ -19,19 +19,10 @@ summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
 class ModifyQueryRequest_v3(BaseModel):
     query_string_list: List[str]
 
-class AnswerQuestionRequest(BaseModel):
-    question: str
-    context: List[str]
-    locations: List[str]
 
 class T5QuestionRequest(BaseModel):
     context: str
 
-
-class AnswerQuestionResponse(BaseModel):
-    answer: str
-    locations: List[str]
-
 class T5Response(BaseModel):
     answer: str
 
@@ -55,21 +46,22 @@ async def modify_query_v3(request: Request):
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error in modifying query v3: {str(e)}")
 
-@app.post("/answer_question", response_model=AnswerQuestionResponse)
-async def answer_question(request: AnswerQuestionRequest):
+@app.post("/answer_question")
+async def answer_question(request: Request):
     try:
+        raw_data = await request.json()
         res_locs = []
         context_string = ''
-        corpus_embeddings = model.encode(request.context, convert_to_tensor=True)
-        query_embeddings = model.encode(request.question, convert_to_tensor=True)
+        corpus_embeddings = model.encode(raw_data['context'], convert_to_tensor=True)
+        query_embeddings = model.encode(raw_data['question'], convert_to_tensor=True)
         hits = util.semantic_search(query_embeddings, corpus_embeddings)
         
         # Collect relevant contexts
         for hit in hits[0]:
             if hit['score'] > 0.4:
                 loc = hit['corpus_id']
-                res_locs.append(request.locations[loc])
-                context_string += request.context[loc] + ' '
+                res_locs.append(raw_data['locations'][loc])
+                context_string += raw_data['context'][loc] + ' '
         
         # If no relevant contexts are found
         if not res_locs:
@@ -77,13 +69,13 @@ async def answer_question(request: AnswerQuestionRequest):
         else:
             # Use the question-answering pipeline
             QA_input = {
-                'question': request.question,
+                'question': raw_data['question'],
                 'context': context_string.replace('\n', ' ')
             }
             result = nlp(QA_input)
             answer = result['answer']
         
-        return AnswerQuestionResponse(answer=answer, locations=res_locs)
+        return JSONResponse(content={'answer':answer, "location":res_locs})
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error in answering question: {str(e)}")