Update app.py

app.py

@@ -1,23 +1,15 @@
-import logging
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from sentence_transformers import SentenceTransformer, util
 from transformers import pipeline
 
-# Set up logging
-logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
-logger = logging.getLogger(__name__)
-
 # Initialize FastAPI app
 app = FastAPI()
 
-# Log model loading
-logger.info("Loading models...")
 # Load models
 model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
 question_model = "deepset/tinyroberta-squad2"
 nlp = pipeline('question-answering', model=question_model, tokenizer=question_model)
-logger.info("Models loaded successfully.")
 
 # Define request models
 class ModifyQueryRequest(BaseModel):
@@ -25,9 +17,10 @@ class ModifyQueryRequest(BaseModel):
 
 class AnswerQuestionRequest(BaseModel):
     question: str
-    context: dict
+    context: list
+    locations: list
 
-# Define response models
+# Define response models (if needed)
 class ModifyQueryResponse(BaseModel):
     embeddings: list
 
@@ -38,50 +31,38 @@ class AnswerQuestionResponse(BaseModel):
 # Define API endpoints
 @app.post("/modify_query", response_model=ModifyQueryResponse)
 async def modify_query(request: ModifyQueryRequest):
-    logger.info(f"Received /modify_query request: {request.query_string}")
     try:
         binary_embeddings = model.encode([request.query_string], precision="binary")
-        logger.info("Embeddings generated successfully.")
         return ModifyQueryResponse(embeddings=binary_embeddings[0].tolist())
     except Exception as e:
-        logger.error(f"Error generating embeddings: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))
 
 @app.post("/answer_question", response_model=AnswerQuestionResponse)
 async def answer_question(request: AnswerQuestionRequest):
-    logger.info(f"Received /answer_question request: {request.question}")
     try:
         res_locs = []
         context_string = ''
-        
-        corpus_embeddings = model.encode(request.context['context'], convert_to_tensor=True)
+        corpus_embeddings = model.encode(request.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'] > 0.5:
+            if hit['score'] > .5:
                 loc = hit['corpus_id']
-                res_locs.append(request.context['locations'][loc])
-                context_string += request.context['context'][loc] + ' '
-
+                res_locs.append(request.locations[loc])
+                context_string += request.context[loc] + ' '
         if len(res_locs) == 0:
             ans = "Sorry, I couldn't find any results for your query."
-            logger.info("No relevant context found.")
         else:
             QA_input = {
                 'question': request.question,
-                'context': context_string.replace('\n', ' ')
+                'context': context_string.replace('\n',' ')
             }
             result = nlp(QA_input)
             ans = result['answer']
-            logger.info("Answer generated successfully.")
-        
-        return AnswerQuestionResponse(answer=ans, locations=res_locs)
+        return AnswerQuestionResponse(answer=ans, locations = res_locs)
     except Exception as e:
-        logger.error(f"Error answering question: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))
 
 if __name__ == "__main__":
     import uvicorn
-    logger.info("Starting FastAPI server...")
     uvicorn.run(app, host="0.0.0.0", port=8000)