Update response_generator.py

response_generator.py

@@ -8,120 +8,214 @@ 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)
-        logger.info(f"Tokenizer loaded from {Config.LLM_MODEL}")
+        """
+        Initialize Arabic LLM for smart generation
+        """
+        # Use AraGPT2 for Arabic text generation
+        model_name = "aubmindlab/aragpt2-base"  # Free Arabic GPT model
+        
+        try:
+            logger.info(f"Loading Arabic generation model: {model_name}")
+            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+            self.model = AutoModelForCausalLM.from_pretrained(model_name)
+            
+            # Set pad token
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+            
+            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            self.model.to(self.device)
+            logger.info(f"Model loaded successfully on {self.device}")
+            self.use_generation = True
+            
+        except Exception as e:
+            logger.warning(f"Could not load generation model: {e}")
+            self.use_generation = False
+            self.tokenizer = None
+            self.model = None
     
     def generate_response(self, query, relevant_docs):
+        """
+        Generate intelligent answer from documents
+        """
         try:
             if len(relevant_docs) == 0:
                 return "عذرًا، لم أجد أي معلومات ذات صلة في المستندات."
             
-            # NEW: Use multi-chunk synthesis for better answers
-            response = self._create_multi_chunk_answer(query, relevant_docs)
-            return response
+            # Use smart generation if model available
+            if self.use_generation and self.model is not None:
+                return self._generate_smart_answer(query, relevant_docs)
+            else:
+                # Fallback: smart extraction
+                return self._generate_extractive_answer(query, relevant_docs)
                 
         except Exception as e:
             logger.error(f"Error generating response: {e}")
             return "عذرًا، لم أتمكن من إنشاء استجابة بسبب خطأ ما."
     
-    def _prepare_context(self, relevant_docs):
-        # 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]
-    
-    # NEW: Multi-chunk answer synthesis
-    def _create_multi_chunk_answer(self, query, relevant_docs):
+    def _generate_smart_answer(self, query, relevant_docs):
         """
-        Create a comprehensive answer by combining top relevant chunks
+        Generate smart answer using LLM
         """
-        # Get top 3 chunks
+        # Get top 3 most relevant chunks
         top_chunks = relevant_docs.head(3)
         
-        # Build answer
-        answer_parts = []
-        answer_parts.append("**بناءً على المستندات المتاحة:**\n")
+        # Combine context
+        contexts = []
+        for idx, row in top_chunks.iterrows():
+            contexts.append(row['content'][:200])  # Limit each chunk to 200 chars
         
-        for idx, (_, row) in enumerate(top_chunks.iterrows(), 1):
-            content = row['content'].strip()
-            similarity = row.get('similarity_score', 0)
-            
-            # Add chunk with similarity indicator
-            if len(top_chunks) > 1:
-                # NEW: Show similarity score for transparency
-                answer_parts.append(f"\n**[{idx}] مطابقة: {similarity:.0%}**")
-            
-            answer_parts.append(content)
+        combined_context = " ".join(contexts)
         
-        # NEW: Add metadata footer
-        answer_parts.append("\n\n---")
+        # Limit total context
+        if len(combined_context) > 400:
+            combined_context = combined_context[:400]
         
-        # NEW: Show number of sources
-        unique_files = top_chunks['path'].nunique()
-        answer_parts.append(f"**عدد المصادر:** {unique_files}")
+        # Create prompt for generation
+        prompt = f"""النص: {combined_context}
+
+السؤال: {query}
+
+الإجابة المختصرة:"""
         
-        # NEW: Show score range for quality indicator
-        if 'similarity_score' in top_chunks.columns:
-            scores = top_chunks['similarity_score'].values
-            if len(scores) == 1:
-                answer_parts.append(f" | **المطابقة:** {scores[0]:.0%}")
+        try:
+            # Tokenize
+            inputs = self.tokenizer(
+                prompt,
+                return_tensors="pt",
+                max_length=512,
+                truncation=True,
+                padding=True
+            ).to(self.device)
+            
+            # Generate answer
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    inputs['input_ids'],
+                    attention_mask=inputs['attention_mask'],
+                    max_new_tokens=100,      # Generate max 100 new tokens
+                    min_length=20,           # Minimum length
+                    num_beams=3,             # Beam search for better quality
+                    no_repeat_ngram_size=3,  # Avoid repetition
+                    do_sample=True,          # Sampling for diversity
+                    top_k=50,
+                    top_p=0.9,
+                    temperature=0.7,
+                    pad_token_id=self.tokenizer.pad_token_id,
+                    eos_token_id=self.tokenizer.eos_token_id
+                )
+            
+            # Decode the answer
+            full_response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            
+            # Extract only the generated answer (remove the prompt)
+            if "الإجابة المختصرة:" in full_response:
+                answer = full_response.split("الإجابة المختصرة:")[-1].strip()
             else:
-                answer_parts.append(f" | **المطابقة:** {scores[-1]:.0%} - {scores[0]:.0%}")
+                answer = full_response[len(prompt):].strip()
+            
+            # Clean up the answer
+            answer = self._clean_answer(answer)
+            
+            # Add metadata
+            scores = top_chunks['similarity_score'].values if 'similarity_score' in top_chunks.columns else []
+            
+            response = f"""**الإجابة:**
+
+{answer}
+
+---
+**عدد المصادر:** {len(top_chunks)} chunks"""
+            
+            if len(scores) > 0:
+                response += f" | **دقة المطابقة:** {scores[-1]:.0%} - {scores[0]:.0%}"
+            
+            logger.info(f"Generated answer length: {len(answer)} chars")
+            
+            return response
+            
+        except Exception as e:
+            logger.error(f"Generation failed: {e}")
+            # Fallback to extractive method
+            return self._generate_extractive_answer(query, relevant_docs)
+    
+    def _clean_answer(self, answer):
+        """
+        Clean and format the generated answer
+        """
+        # Remove incomplete sentences at the end
+        if answer and not answer.endswith(('.', '؟', '!')):
+            # Find last complete sentence
+            last_period = max(answer.rfind('.'), answer.rfind('؟'), answer.rfind('!'))
+            if last_period > 20:  # Keep if reasonable length
+                answer = answer[:last_period + 1]
+        
+        # Remove extra whitespace
+        answer = ' '.join(answer.split())
+        
+        # Limit length
+        if len(answer) > 500:
+            answer = answer[:500]
+            # Try to end at sentence
+            last_period = max(answer.rfind('.'), answer.rfind('؟'))
+            if last_period > 200:
+                answer = answer[:last_period + 1]
         
-        return "\n".join(answer_parts)
+        return answer.strip()
     
-    def _create_extractive_answer(self, query, context, relevant_docs):
+    def _generate_extractive_answer(self, query, relevant_docs):
         """
-        Create an answer by extracting relevant information from documents
+        Smart extraction without generation model (fallback)
         """
-        # Get the most relevant document
-        most_relevant = relevant_docs.iloc[0]['content']
+        top_chunks = relevant_docs.head(3)
         
-        # 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 + "..."
+        # Analyze query to find most relevant chunk
+        query_lower = query.lower()
+        
+        best_chunk = None
+        best_score = 0
         
-        # NEW: Include similarity score if available
-        similarity_info = ""
-        if 'similarity_score' in relevant_docs.columns:
-            score = relevant_docs.iloc[0]['similarity_score']
-            similarity_info = f"\n**المطابقة:** {score:.0%}"
+        for idx, row in top_chunks.iterrows():
+            content = row['content']
+            score = row.get('similarity_score', 0)
+            
+            # Boost score if query keywords appear in content
+            content_lower = content.lower()
+            keyword_matches = sum(1 for word in query_lower.split() if len(word) > 2 and word in content_lower)
+            adjusted_score = score + (keyword_matches * 0.05)
+            
+            if adjusted_score > best_score:
+                best_score = adjusted_score
+                best_chunk = content
         
-        # Format the response
-        response = f"""بناءً على المستندات المتاحة:
+        # Format answer
+        if best_chunk:
+            # Truncate if too long
+            if len(best_chunk) > 400:
+                best_chunk = best_chunk[:400]
+                # Find sentence end
+                last_period = max(best_chunk.rfind('.'), best_chunk.rfind('؟'))
+                if last_period > 100:
+                    best_chunk = best_chunk[:last_period + 1]
+            
+            response = f"""**الإجابة:**
 
-{most_relevant}
+{best_chunk}
 
 ---
-المصدر: {relevant_docs.iloc[0]['path']}{similarity_info}"""
+**المصدر:** {len(top_chunks)} chunks | **المطابقة:** {best_score:.0%}"""
+            
+            return response
         
-        return response
-    
-    def _create_prompt(self, query, context):
-        return f"""مستند قانوني:
-{context}
-
-سؤال:
-{query}
-
-إجابة:"""
+        return "عذرًا، لم أتمكن من إيجاد إجابة مناسبة."
     
     def update_model(self, new_model_name):
         try:
             self.tokenizer = AutoTokenizer.from_pretrained(new_model_name)
+            self.model = AutoModelForCausalLM.from_pretrained(new_model_name)
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+            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}")
-
+            logger.error(f"Error updating model: {e}")