Update main.py

main.py

@@ -136,56 +136,70 @@ import torch
 
 # Model options (ordered by preference)
 QA_MODELS = [
-    "google/flan-t5-small",  # Lightweight default
-    "google/flan-t5-base",   # Medium option
-    "facebook/bart-large-cnn"  # Fallback option
+    {"name": "google/flan-t5-small", "max_length": 512},
+    {"name": "facebook/bart-large-cnn", "max_length": 1024}
 ]
 
-qa_model = None
-current_model_name = None
+class QASystem:
+    def __init__(self):
+        self.model = None
+        self.tokenizer = None
+        self.current_model = None
+        self.device = 0 if torch.cuda.is_available() else -1
 
-def get_qa_model():
-    global qa_model, current_model_name
-    
-    if qa_model is not None:
-        return qa_model
+    def load_model(self):
+        for model_info in QA_MODELS:
+            try:
+                logger.info(f"Loading model: {model_info['name']}")
+                
+                self.tokenizer = AutoTokenizer.from_pretrained(model_info["name"])
+                self.model = AutoModelForSeq2SeqLM.from_pretrained(
+                    model_info["name"],
+                    device_map="auto",
+                    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+                )
+                self.current_model = model_info
+                logger.info(f"Successfully loaded {model_info['name']}")
+                return True
+                
+            except Exception as e:
+                logger.warning(f"Failed to load {model_info['name']}: {str(e)}")
+                continue
+        
+        logger.error("All model loading attempts failed")
+        return False
 
-    # Try each model in order until one works
-    for model_name in QA_MODELS:
+    def generate_answer(self, question: str, context: Optional[str] = None):
         try:
-            logger.info(f"Attempting to load model: {model_name}")
+            if context:
+                input_text = f"question: {question} context: {context[:2000]}"
+            else:
+                input_text = f"question: {question}"
             
-            tokenizer = AutoTokenizer.from_pretrained(model_name)
-            model = AutoModelForSeq2SeqLM.from_pretrained(
-                model_name,
-                device_map="auto" if torch.cuda.is_available() else None,
-                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-            )
+            inputs = self.tokenizer(
+                input_text,
+                return_tensors="pt",
+                truncation=True,
+                max_length=self.current_model["max_length"]
+            ).to(self.device)
             
-            qa_model = pipeline(
-                "text2text-generation",
-                model=model,
-                tokenizer=tokenizer,
-                device=0 if torch.cuda.is_available() else -1
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=200,
+                num_beams=4,
+                early_stopping=True
             )
             
-            current_model_name = model_name
-            logger.info(f"Successfully loaded model: {model_name}")
-            return qa_model
+            return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
             
         except Exception as e:
-            logger.warning(f"Failed to load {model_name}: {str(e)}")
-            continue
+            logger.error(f"Generation failed: {str(e)}")
+            raise
+
+# Initialize QA system
+qa_system = QASystem()
+
 
-    logger.error("All model loading attempts failed")
-    raise HTTPException(
-        status_code=500,
-        detail={
-            "error": "QA system initialization failed",
-            "tried_models": QA_MODELS,
-            "suggestion": "Check available memory or try smaller models"
-        }
-    )
 
 
     
@@ -878,43 +892,63 @@ from typing import Optional
 
 @app.post("/qa")
 async def question_answering(
-    request: Request,
     question: str = Form(...),
     file: Optional[UploadFile] = File(None),
     language: str = Form("en")
 ):
+    # Initialize model if not loaded
+    if not qa_system.model:
+        if not qa_system.load_model():
+            raise HTTPException(
+                500,
+                detail={
+                    "error": "System initialization failed",
+                    "tried_models": [m["name"] for m in QA_MODELS],
+                    "suggestion": "Check logs for loading errors"
+                }
+            )
+
     try:
-        # Initialize model (with fallback)
+        # Process file if provided
+        context = None
+        if file:
+            try:
+                file_ext, content = await process_uploaded_file(file)
+                context = extract_text(content, file_ext)
+                context = re.sub(r'\s+', ' ', context).strip()[:3000]
+            except Exception as e:
+                logger.error(f"File processing failed: {str(e)}")
+                raise HTTPException(422, detail=f"File processing error: {str(e)}")
+
+        # Generate answer
         try:
-            qa_pipeline = get_qa_model()
+            answer = qa_system.generate_answer(question, context)
+            
+            return {
+                "question": question,
+                "answer": answer,
+                "model": qa_system.current_model["name"],
+                "source": "document" if context else "general",
+                "language": language
+            }
+            
         except Exception as e:
-            logger.critical(f"Model loading failed: {str(e)}")
-            raise HTTPException(500, "Could not initialize any QA model")
-
-        # Rest of your existing endpoint logic...
-        # [Keep all your existing file processing and QA code]
-        
-        return {
-            "question": question,
-            "answer": result[0]["generated_text"],
-            "model_used": current_model_name,  # Add this to responses
-            "source": "document" if file else "general knowledge",
-            "language": language
-        }
+            logger.error(f"Answer generation failed: {str(e)}")
+            raise HTTPException(
+                500,
+                detail={
+                    "error": "Answer generation failed",
+                    "model": qa_system.current_model["name"],
+                    "input_length": len(question) + (len(context) if context else 0),
+                    "suggestion": "Try simplifying your question or reducing document size"
+                }
+            )
 
     except HTTPException:
         raise
     except Exception as e:
-        logger.error(f"QA processing failed: {str(e)}")
-        raise HTTPException(
-            500,
-            detail={
-                "error": "QA processing failed",
-                "model": current_model_name,
-                "input_question": question[:100] + "..." if question else None,
-                "file_type": file.filename.split('.')[-1] if file else None
-            }
-        )
+        logger.critical(f"Unexpected error: {str(e)}")
+        raise HTTPException(500, "Internal server error")