Update app.py

app.py

@@ -1,6 +1,6 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from sentence_transformers import SentenceTransformer
+from sentence_transformers import SentenceTransformer, utils
 from transformers import pipeline
 
 # Initialize FastAPI app
@@ -17,7 +17,7 @@ class ModifyQueryRequest(BaseModel):
 
 class AnswerQuestionRequest(BaseModel):
     question: str
-    context: str
+    context: dict
 
 # Define response models (if needed)
 class ModifyQueryResponse(BaseModel):
@@ -25,6 +25,7 @@ class ModifyQueryResponse(BaseModel):
 
 class AnswerQuestionResponse(BaseModel):
     answer: str
+    locations: list
 
 # Define API endpoints
 @app.post("/modify_query", response_model=ModifyQueryResponse)
@@ -38,12 +39,26 @@ async def modify_query(request: ModifyQueryRequest):
 @app.post("/answer_question", response_model=AnswerQuestionResponse)
 async def answer_question(request: AnswerQuestionRequest):
     try:
-        QA_input = {
-            'question': request.question,
-            'context': request.context
-        }
-        result = nlp(QA_input)
-        return AnswerQuestionResponse(answer=result['answer'])
+        res_locs = []
+        context_string = ''
+        corpus_embeddings = model.encode(request.context['context'], convert_to_tensor=True)
+        query_embeddings = model.encode(request.question, convert_to_tensor=True)
+        hits = util.semantic_search(query_embeddings, corpus_embeddings)
+        for hit in hits:
+            if hit['score'] > .5:
+                loc = hit['corpus_id']
+                res_locs.append(request.context['locations'][loc])
+                context_string += request.context['context'][loc] + ' '
+        if len(res_locs) == 0:
+            ans = "Sorry, I couldn't find any results for your query."
+        else:
+            QA_input = {
+                'question': request.question,
+                'context': context_string.replace.('\n',' ')
+            }
+            result = nlp(QA_input)
+            ans = result['answer']
+        return AnswerQuestionResponse(answer=ans, locations = res_locs)
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))