models/text_generator.py

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import json
import random
from typing import Dict, Any, List

class QwenTextGenerator:
    """Text generation using Qwen2.5-0.5B-Instruct for monster traits and dialogue"""
    
    def __init__(self, device: str = "cuda"):
        self.device = device if torch.cuda.is_available() else "cpu"
        self.model = None
        self.tokenizer = None
        self.model_id = "Qwen/Qwen2.5-0.5B-Instruct"
        
        # Generation parameters
        self.max_new_tokens = 150
        self.temperature = 0.8
        self.top_p = 0.9
        
        # Monster trait templates
        self.trait_categories = {
            'elements': ['fire', 'water', 'earth', 'wind', 'electric', 'ice', 'nature', 'dark', 'light', 'neutral'],
            'personalities': ['brave', 'timid', 'aggressive', 'gentle', 'playful', 'serious', 'loyal', 'independent', 'curious', 'protective'],
            'body_types': ['bipedal', 'quadruped', 'serpentine', 'avian', 'aquatic', 'insectoid', 'humanoid', 'amorphous'],
            'sizes': ['tiny', 'small', 'medium', 'large', 'giant'],
            'special_features': ['wings', 'horns', 'tail', 'spikes', 'fur', 'scales', 'armor', 'crystals', 'flames', 'aura']
        }
    
    def load_model(self):
        """Lazy load the text generation model"""
        if self.model is None:
            try:
                # Load tokenizer
                self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
                
                # Model configuration
                torch_dtype = torch.float16 if self.device == "cuda" else torch.float32
                
                self.model = AutoModelForCausalLM.from_pretrained(
                    self.model_id,
                    torch_dtype=torch_dtype,
                    device_map="auto" if self.device == "cuda" else None,
                    low_cpu_mem_usage=True
                )
                
                if self.device == "cpu":
                    self.model.to(self.device)
                
            except Exception as e:
                print(f"Failed to load text generation model: {e}")
                raise
    
    def generate_traits(self, description: str) -> Dict[str, Any]:
        """Generate monster traits from description"""
        try:
            self.load_model()
            
            # Create prompt for trait generation
            prompt = self._create_trait_prompt(description)
            
            # Generate response
            inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
            
            with torch.no_grad():
                outputs = self.model.generate(
                    **inputs,
                    max_new_tokens=self.max_new_tokens,
                    temperature=self.temperature,
                    top_p=self.top_p,
                    do_sample=True,
                    pad_token_id=self.tokenizer.eos_token_id
                )
            
            response = self.tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
            
            # Parse traits from response
            traits = self._parse_traits(response, description)
            
            return traits
            
        except Exception as e:
            print(f"Error generating traits: {e}")
            return self._generate_fallback_traits(description)
    
    def generate_dialogue(self, traits: Dict[str, Any]) -> str:
        """Generate monster dialogue (emoji + numbers)"""
        try:
            # Create emoji dialogue based on personality and mood
            personality = traits.get('personality', 'neutral')
            
            # Emoji mapping for personalities
            emoji_map = {
                'brave': ['💪', '🔥', '⚔️', '🛡️'],
                'timid': ['😰', '🥺', '💦', '❓'],
                'aggressive': ['😤', '💢', '🔥', '⚡'],
                'gentle': ['💚', '🌸', '✨', '🌟'],
                'playful': ['😊', '🎮', '🎯', '🎪'],
                'serious': ['🤖', '📊', '⚡', '💯'],
                'loyal': ['💖', '🤝', '🛡️', '⭐'],
                'independent': ['🚀', '🌍', '🔮', '💫'],
                'curious': ['🔍', '❓', '💡', '🌟'],
                'protective': ['🛡️', '💪', '🏰', '⚔️']
            }
            
            # Get appropriate emojis
            emojis = emoji_map.get(personality, ['🤖', '💚', '✨'])
            selected_emojis = random.sample(emojis, min(2, len(emojis)))
            
            # Generate status numbers (representing monster's current state)
            hp_percent = random.randint(70, 100)
            happiness = random.randint(60, 95)
            energy = random.randint(50, 90)
            
            # Create dialogue
            dialogue = f"{selected_emojis[0]}{selected_emojis[1] if len(selected_emojis) > 1 else '💚'}"
            dialogue += f"{hp_percent}️⃣{happiness}️⃣"
            
            return dialogue
            
        except Exception as e:
            print(f"Error generating dialogue: {e}")
            return "🤖💚9️⃣0️⃣"
    
    def _create_trait_prompt(self, description: str) -> str:
        """Create prompt for trait generation"""
        prompt = f"""<|im_start|>system
You are a creative game designer creating unique digital monsters. Generate detailed traits for a monster based on the description.
<|im_end|>
<|im_start|>user
Create traits for this monster: {description}

Include: name, species, element, personality, appearance details, and special abilities.
<|im_end|>
<|im_start|>assistant
"""
        return prompt
    
    def _parse_traits(self, response: str, original_description: str) -> Dict[str, Any]:
        """Parse traits from model response"""
        traits = {
            'description': original_description,
            'raw_response': response
        }
        
        # Extract name
        if "name:" in response.lower():
            name_start = response.lower().find("name:") + 5
            name_end = response.find("\n", name_start)
            if name_end == -1:
                name_end = len(response)
            traits['name'] = response[name_start:name_end].strip()
        else:
            traits['name'] = self._generate_name()
        
        # Extract or assign element
        element_found = False
        for element in self.trait_categories['elements']:
            if element in response.lower():
                traits['element'] = element
                element_found = True
                break
        
        if not element_found:
            traits['element'] = random.choice(self.trait_categories['elements'])
        
        # Extract or assign personality
        personality_found = False
        for personality in self.trait_categories['personalities']:
            if personality in response.lower():
                traits['personality'] = personality
                personality_found = True
                break
        
        if not personality_found:
            traits['personality'] = random.choice(self.trait_categories['personalities'])
        
        # Extract appearance
        traits['appearance'] = self._extract_appearance(response)
        
        # Extract abilities
        traits['abilities'] = self._extract_abilities(response, traits['element'])
        
        # Add color scheme based on element
        traits['color_scheme'] = self._get_color_scheme(traits['element'])
        
        return traits
    
    def _generate_name(self) -> str:
        """Generate a random monster name"""
        prefixes = ['Pyro', 'Aqua', 'Terra', 'Aero', 'Volt', 'Cryo', 'Flora', 'Shadow', 'Lumi', 'Neo']
        suffixes = ['mon', 'beast', 'guard', 'wing', 'claw', 'fang', 'horn', 'tail', 'byte', 'spark']
        
        return random.choice(prefixes) + random.choice(suffixes)
    
    def _extract_appearance(self, response: str) -> str:
        """Extract appearance description"""
        appearance_keywords = ['appearance', 'looks like', 'resembles', 'body', 'color', 'size']
        
        for keyword in appearance_keywords:
            if keyword in response.lower():
                start = response.lower().find(keyword)
                end = response.find('.', start)
                if end == -1:
                    end = response.find('\n', start)
                if end == -1:
                    end = len(response)
                
                return response[start:end].strip()
        
        # Fallback appearance
        body_type = random.choice(self.trait_categories['body_types'])
        size = random.choice(self.trait_categories['sizes'])
        feature = random.choice(self.trait_categories['special_features'])
        
        return f"A {size} {body_type} creature with {feature}"
    
    def _extract_abilities(self, response: str, element: str) -> List[str]:
        """Extract or generate abilities"""
        abilities = []
        
        ability_keywords = ['ability', 'power', 'skill', 'can', 'capable']
        for keyword in ability_keywords:
            if keyword in response.lower():
                # Try to extract abilities from response
                start = response.lower().find(keyword)
                end = response.find('.', start)
                if end > start:
                    ability_text = response[start:end]
                    abilities.append(ability_text.strip())
        
        # If no abilities found, generate based on element
        if not abilities:
            element_abilities = {
                'fire': ['Flame Burst', 'Heat Wave', 'Ember Shield'],
                'water': ['Aqua Jet', 'Bubble Shield', 'Tidal Wave'],
                'earth': ['Rock Throw', 'Earthquake', 'Stone Armor'],
                'wind': ['Gust', 'Tornado', 'Wind Shield'],
                'electric': ['Thunder Shock', 'Static Field', 'Lightning Speed'],
                'ice': ['Ice Beam', 'Frost Armor', 'Blizzard'],
                'nature': ['Vine Whip', 'Healing Bloom', 'Nature\'s Guard'],
                'dark': ['Shadow Strike', 'Dark Pulse', 'Void Shield'],
                'light': ['Light Beam', 'Healing Light', 'Radiant Shield'],
                'neutral': ['Tackle', 'Defense Curl', 'Focus']
            }
            
            abilities = random.sample(
                element_abilities.get(element, element_abilities['neutral']), 
                2
            )
        
        return abilities
    
    def _get_color_scheme(self, element: str) -> str:
        """Get color scheme based on element"""
        color_schemes = {
            'fire': 'red and orange with yellow accents',
            'water': 'blue and cyan with white highlights',
            'earth': 'brown and green with stone textures',
            'wind': 'white and light blue with swirling patterns',
            'electric': 'yellow and blue with sparking effects',
            'ice': 'light blue and white with crystalline features',
            'nature': 'green and brown with leaf patterns',
            'dark': 'black and purple with shadow effects',
            'light': 'white and gold with glowing aura',
            'neutral': 'gray and silver with balanced tones'
        }
        
        return color_schemes.get(element, 'varied colors with unique patterns')
    
    def _generate_fallback_traits(self, description: str) -> Dict[str, Any]:
        """Generate fallback traits if model fails"""
        element = random.choice(self.trait_categories['elements'])
        personality = random.choice(self.trait_categories['personalities'])
        
        return {
            'name': self._generate_name(),
            'species': 'Digital Monster',
            'element': element,
            'personality': personality,
            'appearance': f"A unique {random.choice(self.trait_categories['sizes'])} digital creature",
            'color_scheme': self._get_color_scheme(element),
            'abilities': self._extract_abilities("", element),
            'description': description
        }
    
    def to(self, device: str):
        """Move model to specified device"""
        self.device = device
        if self.model:
            self.model.to(device)
    
    def __del__(self):
        """Cleanup when object is destroyed"""
        if self.model:
            del self.model
        if self.tokenizer:
            del self.tokenizer
        torch.cuda.empty_cache()