Update response_generator.py

response_generator.py

@@ -1,6 +1,6 @@
 # response_generator.py
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from utils import setup_logger
 from config import Config
 
@@ -8,47 +8,65 @@ logger = setup_logger('response_generator')
 
 class ResponseGenerator:
     def __init__(self):
+        # Use a simpler approach with a summarization/QA pipeline
+        # Since BERT-based models don't generate text, we'll create a simple retrieval-based response
         self.tokenizer = AutoTokenizer.from_pretrained(Config.LLM_MODEL)
-        self.model = AutoModelForCausalLM.from_pretrained(Config.LLM_MODEL)
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.model.to(self.device)
-        logger.info(f"Model loaded and moved to {self.device}")
-
+        logger.info(f"Tokenizer loaded from {Config.LLM_MODEL}")
+    
     def generate_response(self, query, relevant_docs):
         try:
-            context = self._prepare_context(relevant_docs)
-            prompt = self._create_prompt(query, context)
-            input_ids = self.tokenizer.encode(prompt, return_tensors="pt").to(self.device)
+            if len(relevant_docs) == 0:
+                return "عذرًا، لم أجد أي معلومات ذات صلة في المستندات."
             
-            attention_mask = input_ids.ne(self.tokenizer.pad_token_id).float()
-            
-            with torch.no_grad():
-                output = self.model.generate(
-                    input_ids,
-                    attention_mask=attention_mask,
-                    max_length=Config.MAX_LENGTH,
-                    num_return_sequences=1,
-                    no_repeat_ngram_size=2,
-                    do_sample=True,
-                    top_k=50,
-                    top_p=0.95,
-                    temperature=0.7
-                )
+            # Get the most relevant document (first one)
+            context = self._prepare_context(relevant_docs)
             
-            response = self.tokenizer.decode(output[0], skip_special_tokens=True)
-            return self._extract_answer(response)
-        
+            # For BERT-based models, we do extractive QA instead of generation
+            # Return the most relevant context as the answer
+            response = self._create_extractive_answer(query, context, relevant_docs)
+            return response
+                
         except Exception as e:
             logger.error(f"Error generating response: {e}")
-            return "عذرًا، لم أتمكن من إنشاء استجابة بسبب خطأ ما."  # "Sorry, I couldn't generate a response due to an error."
-
+            return "عذرًا، لم أتمكن من إنشاء استجابة بسبب خطأ ما."
+    
     def _prepare_context(self, relevant_docs):
-        # Combine content from relevant documents
-        combined_content = "\n".join(relevant_docs['content'].tolist())
-        # Truncate if too long
-        max_context_length = Config.MAX_LENGTH // 2  # Use half of max_length for context
+        # Take only the top 3 most relevant documents to avoid token limit
+        top_docs = relevant_docs.head(3)
+        combined_content = "\n\n".join(top_docs['content'].tolist())
+        
+        # Limit to 300 characters to stay within token limits
+        max_context_length = 300
         return combined_content[:max_context_length]
+    
+    def _create_extractive_answer(self, query, context, relevant_docs):
+        """
+        Create an answer by extracting relevant information from documents
+        """
+        # Get the most relevant document
+        most_relevant = relevant_docs.iloc[0]['content']
+        
+        # Truncate to reasonable length
+        max_length = 500
+        if len(most_relevant) > max_length:
+            # Try to find a good sentence break
+            truncated = most_relevant[:max_length]
+            last_period = truncated.rfind('.')
+            if last_period > 0:
+                most_relevant = truncated[:last_period + 1]
+            else:
+                most_relevant = truncated + "..."
+        
+        # Format the response
+        response = f"""بناءً على المستندات المتاحة:
+
+{most_relevant}
 
+---
+المصدر: {relevant_docs.iloc[0]['path']}"""
+        
+        return response
+    
     def _create_prompt(self, query, context):
         return f"""مستند قانوني:
 {context}
@@ -57,17 +75,10 @@ class ResponseGenerator:
 {query}
 
 إجابة:"""
-
-    def _extract_answer(self, response):
-        # Extract the generated answer from the full response
-        answer_start = response.find("إجابة:") + len("إجابة:")
-        return response[answer_start:].strip()
-
+    
     def update_model(self, new_model_name):
         try:
             self.tokenizer = AutoTokenizer.from_pretrained(new_model_name)
-            self.model = AutoModelForCausalLM.from_pretrained(new_model_name)
-            self.model.to(self.device)
             logger.info(f"Model updated to {new_model_name}")
         except Exception as e:
             logger.error(f"Error updating model: {e}")