Update app.py

app.py

@@ -96,8 +96,50 @@ class DocumentRAG:
         if not context:
             return "No relevant information found in the documents."
         
-        # Improved keyword matching approach
-        query_words = set(query.lower().split())
+        query_lower = query.lower()
+        
+        # Handle "who is" questions specifically
+        if "who is" in query_lower:
+            # Extract name from query
+            name_part = query_lower.replace("who is", "").strip()
+            
+            # Look for professional information in context
+            lines = context.split('\n')
+            name_info = []
+            professional_info = []
+            
+            for line in lines:
+                line = line.strip()
+                if not line or line.startswith('---'):
+                    continue
+                    
+                line_lower = line.lower()
+                
+                # Look for job titles, companies, roles
+                if any(keyword in line_lower for keyword in [
+                    'scientist', 'engineer', 'analyst', 'developer', 'manager', 
+                    'consultant', 'specialist', 'coordinator', 'associate', 'intern',
+                    'at ', 'working at', 'employed', 'position', 'role'
+                ]):
+                    professional_info.append(line)
+                
+                # Look for name and basic info
+                elif any(keyword in line_lower for keyword in [
+                    'name', 'email', 'phone', 'linkedin', 'github', 'experience'
+                ]):
+                    name_info.append(line)
+            
+            # Construct answer
+            if professional_info:
+                answer = f"Based on the resume, {name_part} is " + professional_info[0]
+                if len(professional_info) > 1:
+                    answer += f" and also {professional_info[1]}"
+                return answer
+            elif name_info:
+                return f"The document shows information about {name_part}: " + "; ".join(name_info[:2])
+        
+        # For other questions, use improved keyword matching
+        query_words = set(query_lower.split())
         context_sentences = context.split('.')
         
         # Find sentences that contain query keywords
@@ -115,7 +157,7 @@ class DocumentRAG:
         
         if relevant_sentences:
             # Return the most relevant sentences
-            return '. '.join(relevant_sentences[:3]) + '.'
+            return '. '.join(relevant_sentences[:2]) + '.'
         else:
             # If no exact matches, return first few sentences of context
             first_sentences = context_sentences[:2]
@@ -303,25 +345,27 @@ class DocumentRAG:
             is_mistral = 'mistral' in model_name
             
             if is_mistral:
-                # Improved prompt for Mistral - more flexible
-                prompt = f"""<s>[INST] You are a helpful document assistant. Answer the question based on the provided context. If the exact answer isn't in the context, provide the most relevant information available.
+                # Improved prompt for Mistral with specific instructions
+                prompt = f"""<s>[INST] You are a helpful assistant that answers questions about people based on their resume/document information. 
+
+Answer the question clearly and concisely. For "who is" questions, provide a brief professional summary.
 
-Context:
+Context from document:
 {context[:1500]}
 
 Question: {query}
 
-Please provide a helpful answer based on the available information. [/INST]"""
+Provide a clear, direct answer in 1-2 sentences. [/INST]"""
             else:
                 # Improved prompt for fallback models
-                prompt = f"""Based on the following information, please answer the question:
+                prompt = f"""Answer the question about the person based on their resume information:
 
-Context:
+Resume Information:
 {context[:1000]}
 
 Question: {query}
 
-Answer:"""
+Answer (be direct and concise):"""
 
             # Tokenize with proper handling
             inputs = self.tokenizer(
@@ -336,17 +380,17 @@ Answer:"""
             if torch.cuda.is_available() and next(self.model.parameters()).is_cuda:
                 inputs = {k: v.cuda() for k, v in inputs.items()}
             
-            # Generate with more flexible parameters
+            # Generate with more focused parameters
             with torch.no_grad():
                 outputs = self.model.generate(
                     **inputs,
-                    max_new_tokens=150,
-                    temperature=0.3,      # Slightly higher for more natural responses
+                    max_new_tokens=100,   # Shorter for more focused answers
+                    temperature=0.2,      # Lower for more deterministic responses
                     do_sample=True,
-                    top_p=0.9,
-                    num_beams=2,
+                    top_p=0.8,
+                    num_beams=3,
                     early_stopping=True,
-                    repetition_penalty=1.1,
+                    repetition_penalty=1.2,
                     pad_token_id=self.tokenizer.pad_token_id,
                     eos_token_id=self.tokenizer.eos_token_id
                 )
@@ -359,7 +403,9 @@ Answer:"""
                 answer = full_response.split("[/INST]")[-1].strip()
             else:
                 # For other models, remove the prompt
-                if "Answer:" in full_response:
+                if "Answer (be direct and concise):" in full_response:
+                    answer = full_response.split("Answer (be direct and concise):")[-1].strip()
+                elif "Answer:" in full_response:
                     answer = full_response.split("Answer:")[-1].strip()
                 else:
                     answer = full_response[len(prompt):].strip()
@@ -374,28 +420,45 @@ Answer:"""
             return self.simple_context_answer(query, context)
     
     def clean_answer(self, answer: str) -> str:
-        """Clean up the generated answer"""
+        """Clean up the generated answer with better formatting"""
         if not answer or len(answer) < 5:
             return ""
         
-        # Remove obvious problematic patterns
-        lines = answer.split('\n')
-        cleaned_lines = []
+        # Remove file markers and cleanup
+        answer = answer.replace('--- ', '').replace(' ---', '')
+        answer = answer.replace('.pdf', '').replace('.docx', '').replace('.txt', '')
         
-        for line in lines:
-            line = line.strip()
-            if line and not any(pattern in line.lower() for pattern in [
+        # Split into sentences and clean each
+        sentences = answer.split('.')
+        cleaned_sentences = []
+        
+        for sentence in sentences:
+            sentence = sentence.strip()
+            if not sentence:
+                continue
+                
+            # Skip problematic patterns
+            if any(pattern in sentence.lower() for pattern in [
                 'what are you doing', 'what do you think', 'how are you',
-                'i am an ai', 'i cannot', 'i don\'t know'
+                'i am an ai', 'i cannot', 'i don\'t know', 'linkedin: www',
+                'github:', 'email:', 'mobile:', '+91-'
             ]):
-                cleaned_lines.append(line)
+                continue
+            
+            # Clean up common formatting issues
+            sentence = sentence.replace('  ', ' ')
+            if sentence and len(sentence) > 3:
+                cleaned_sentences.append(sentence)
         
-        cleaned_answer = ' '.join(cleaned_lines)
+        if not cleaned_sentences:
+            return ""
+            
+        # Reconstruct answer
+        cleaned_answer = '. '.join(cleaned_sentences[:2])  # Limit to 2 sentences
         
-        # Limit length to prevent rambling
-        if len(cleaned_answer) > 500:
-            sentences = cleaned_answer.split('.')
-            cleaned_answer = '. '.join(sentences[:3]) + '.'
+        # Add period if missing
+        if cleaned_answer and not cleaned_answer.endswith('.'):
+            cleaned_answer += '.'
         
         return cleaned_answer.strip()