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()