explanation_generator.py

"""
Explanation Generator Module

This module handles the generation of explanations for resume rankings
using the QwQ-32B model from Hugging Face.
"""

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import os
import re

class ExplanationGenerator:
    def __init__(self, model_name="Qwen/QwQ-32B"):
        """Initialize the explanation generator with the specified model"""
        self.model_name = model_name
        self.model = None
        self.tokenizer = None
        self.initialized = False
        
    def load_model(self):
        """Load the model and tokenizer if not already loaded"""
        if not self.initialized:
            try:
                # Check if we have enough VRAM for loading the model
                if torch.cuda.is_available():
                    gpu_memory = torch.cuda.get_device_properties(0).total_memory
                    # QwQ-32B requires at least 32GB VRAM for full precision
                    if gpu_memory >= 32 * (1024**3):  # 32 GB
                        device = "cuda"
                    else:
                        device = "cpu"
                else:
                    device = "cpu"
                
                # Load tokenizer
                self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
                
                # Load model based on available resources
                if device == "cuda":
                    self.model = AutoModelForCausalLM.from_pretrained(
                        self.model_name,
                        torch_dtype=torch.bfloat16,
                        device_map="auto"
                    )
                else:
                    # Fall back to a simpler template-based solution if we can't load the model
                    self.model = None
                    print("Warning: Loading QwQ-32B on CPU is not recommended. Using template-based explanations instead.")
                
                self.initialized = True
            except Exception as e:
                print(f"Error loading QwQ-32B model: {str(e)}")
                print("Falling back to template-based explanations.")
                self.model = None
                self.initialized = True
    
    def generate_explanation(self, resume_text, job_description, score, semantic_score, keyword_score, skills):
        """Generate explanation for why a resume was ranked highly"""
        # Check if we need to load the model
        if not self.initialized:
            self.load_model()
        
        # If the model is loaded and available, use it for generating explanations
        if self.model is not None:
            try:
                # Prepare prompt for QwQ-32B
                prompt = self._create_prompt(resume_text, job_description, score, semantic_score, keyword_score, skills)
                
                # Create messages for chat format
                messages = [
                    {"role": "user", "content": prompt}
                ]
                
                # Apply chat template
                text = self.tokenizer.apply_chat_template(
                    messages,
                    tokenize=False,
                    add_generation_prompt=True
                )
                
                # Tokenize
                inputs = self.tokenizer(text, return_tensors="pt").to(self.model.device)
                
                # Generate response
                output_ids = self.model.generate(
                    **inputs,
                    max_new_tokens=300,
                    temperature=0.6,
                    top_p=0.95,
                    top_k=30
                )
                
                # Decode the response
                response = self.tokenizer.decode(output_ids[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
                
                # Clean up the response
                cleaned_response = self._clean_response(response)
                
                return cleaned_response
                
            except Exception as e:
                print(f"Error generating explanation with QwQ-32B: {str(e)}")
                # Fall back to template-based explanation
                return self._generate_template_explanation(score, semantic_score, keyword_score, skills)
        else:
            # Use template-based explanation if model is not available
            return self._generate_template_explanation(score, semantic_score, keyword_score, skills)
    
    def _create_prompt(self, resume_text, job_description, score, semantic_score, keyword_score, skills):
        """Create a prompt for the explanation generation"""
        # Use only the first 1000 characters of the resume to keep prompt size manageable
        resume_excerpt = resume_text[:1000] + "..." if len(resume_text) > 1000 else resume_text
        
        prompt = f"""You are an AI assistant helping a recruiter understand why a candidate's resume was matched with a job posting.

The resume has been assigned the following scores:
- Overall Match Score: {score:.2f} out of 1.0
- Semantic Relevance Score: {semantic_score:.2f} out of 1.0
- Keyword Match Score: {keyword_score:.2f} out of 1.0

The job description is:
```
{job_description}
```

Based on analysis, the resume contains these skills relevant to the job: {', '.join(skills)}

Resume excerpt:
```
{resume_excerpt}
```

Please provide a short explanation (3-5 sentences) of why this resume received these scores and how well it matches the job requirements. Focus on the relationship between the candidate's experience and the job requirements."""

        return prompt
    
    def _clean_response(self, response):
        """Clean the response from the model"""
        # Remove any thinking or internal processing tokens
        response = re.sub(r'<think>.*?</think>', '', response, flags=re.DOTALL)
        
        # Limit to a reasonable length
        if len(response) > 500:
            sentences = response.split('.')
            shortened = '.'.join(sentences[:5]) + '.'
            return shortened
        
        return response
    
    def _generate_template_explanation(self, score, semantic_score, keyword_score, skills):
        """Generate a template-based explanation when the model is not available"""
        # Simple template-based explanation
        if score > 0.8:
            quality = "excellent"
        elif score > 0.6:
            quality = "good"
        elif score > 0.4:
            quality = "moderate"
        else:
            quality = "limited"
            
        explanation = f"This resume shows {quality} alignment with the job requirements, with an overall score of {score:.2f}. "
        
        if semantic_score > keyword_score:
            explanation += f"The candidate's experience demonstrates strong semantic relevance ({semantic_score:.2f}) to the position, though specific keyword matches ({keyword_score:.2f}) could be improved. "
        else:
            explanation += f"The resume contains many relevant keywords ({keyword_score:.2f}), but could benefit from better contextual alignment ({semantic_score:.2f}) with the job requirements. "
        
        if skills:
            if len(skills) > 3:
                explanation += f"Key skills identified include {', '.join(skills[:3])}, and {len(skills)-3} others that match the job requirements."
            else:
                explanation += f"Key skills identified include {', '.join(skills)}."
        else:
            explanation += "No specific skills were identified that directly match the requirements."
            
        return explanation