Update app.py

app.py

@@ -1,8 +1,9 @@
 import os
 import pickle
 import numpy as np
-from flask import Flask, request, jsonify
-from flask_cors import CORS
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
 from transformers import (
     AutoTokenizer, 
     AutoModelForSeq2SeqLM, 
@@ -17,13 +18,26 @@ import nltk
 import torch
 import pandas as pd
 
-app = Flask(__name__)
-CORS(app)
+# Initialize FastAPI app
+app = FastAPI()
+
+# Add CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
 # Global variables for models and data
 models = {}
 data = {}
 
+class QueryRequest(BaseModel):
+    query: str
+    language_code: int = 0
+
 def init_nltk():
     """Initialize NLTK resources"""
     try:
@@ -78,7 +92,6 @@ def load_embeddings():
             print(f"Error: {embeddings_path} not found")
             return False
             
-        # Custom unpickler to handle potential compatibility issues
         class CustomUnpickler(pickle.Unpickler):
             def find_class(self, module, name):
                 if module == "__main__":
@@ -193,11 +206,9 @@ def generate_answer(query, context, max_length=860, temperature=0.2):
         
         response = models['llm_tokenizer'].decode(outputs[0], skip_special_tokens=True)
         
-        # Clean up the response
         if "Answer:" in response:
             response = response.split("Answer:")[-1].strip()
         
-        # Remove incomplete sentences at the end
         sentences = nltk.sent_tokenize(response)
         if sentences:
             return " ".join(sentences)
@@ -247,8 +258,8 @@ def rerank_documents(query, doc_texts):
         print(f"Error reranking documents: {e}")
         return np.zeros(len(doc_texts))
 
-@app.route('/health', methods=['GET'])
-def health_check():
+@app.get("/health")
+async def health_check():
     """Health check endpoint"""
     status = {
         'status': 'healthy',
@@ -256,95 +267,79 @@ def health_check():
         'embeddings_loaded': bool(data.get('embeddings')),
         'documents_loaded': not data.get('df', pd.DataFrame()).empty
     }
-    return jsonify(status)
+    return status
 
-@app.route('/api/query', methods=['POST'])
-def process_query():
+@app.post("/api/query")
+async def process_query(request: QueryRequest):
     """Main query processing endpoint"""
     try:
-        if not request.is_json:
-            return jsonify({'error': 'Request must be JSON', 'success': False}), 400
-            
-        data = request.json
-        if not data or 'query' not in data:
-            return jsonify({'error': 'No query provided', 'success': False}), 400
+        query_text = request.query
+        language_code = request.language_code
 
-        query_text = data['query']
-        language_code = data.get('language_code', 0)
-
-        # Basic response if no models or data are loaded
         if not models or not data.get('embeddings'):
-            return jsonify({
-                'answer': 'The system is currently initializing. Please try again in a few minutes.',
-                'success': False
-            }), 503
+            raise HTTPException(
+                status_code=503,
+                detail="The system is currently initializing. Please try again in a few minutes."
+            )
 
-        # Process query with available models and data
         try:
-            # Handle Arabic queries
             if language_code == 0:
                 query_text = translate_text(query_text, 'ar_to_en')
 
-            # Get query embedding and find relevant documents
             query_embedding = models['embedding'].encode([query_text])
             relevant_docs = query_embeddings(query_embedding)
             
             if not relevant_docs:
-                return jsonify({
+                return {
                     'answer': 'No relevant information found. Please try a different query.',
                     'success': True
-                })
+                }
 
-            # Retrieve and process documents
             doc_texts = [retrieve_document_text(doc_id) for doc_id, _ in relevant_docs]
             doc_texts = [text for text in doc_texts if text.strip()]
             
             if not doc_texts:
-                return jsonify({
+                return {
                     'answer': 'Unable to retrieve relevant documents. Please try again.',
                     'success': True
-                })
+                }
 
-            # Rerank documents
             rerank_scores = rerank_documents(query_text, doc_texts)
             ranked_texts = [text for _, text in sorted(zip(rerank_scores, doc_texts), reverse=True)]
             
-            # Combine top documents
             context = " ".join(ranked_texts[:3])
-            
-            # Generate answer
             answer = generate_answer(query_text, context)
             
-            # Translate answer back to Arabic if needed
             if language_code == 0:
                 answer = translate_text(answer, 'en_to_ar')
 
-            return jsonify({
+            return {
                 'answer': answer,
                 'success': True
-            })
+            }
 
         except Exception as e:
             print(f"Error processing query: {e}")
-            return jsonify({
-                'error': 'An error occurred while processing your query',
-                'success': False
-            }), 500
+            raise HTTPException(
+                status_code=500,
+                detail="An error occurred while processing your query"
+            )
 
     except Exception as e:
         print(f"Error in process_query: {e}")
-        return jsonify({
-            'error': str(e),
-            'success': False
-        }), 500
+        raise HTTPException(
+            status_code=500,
+            detail=str(e)
+        )
 
-# Initialize everything when the app starts
+# Initialize application
 print("Initializing application...")
 init_success = init_nltk() and load_models() and load_data()
 
 if not init_success:
     print("Warning: Application initialized with partial functionality")
 
+# For running locally
 if __name__ == "__main__":
-    app.run(host='0.0.0.0', port=7860)
-    
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)