Update RockPaperScissor/services/LLM_service.py

RockPaperScissor/services/LLM_service.py

@@ -4,59 +4,58 @@ from typing import List, Dict, Optional
 import re
 import os
 
+# Better handling of spaces GPU decorator
+try:
+    from spaces import GPU
+    print("[LLMService] spaces.GPU imported successfully")
+except ImportError:
+    print("[LLMService] spaces.GPU not available, using fallback")
+    def GPU(f): 
+        return f  # fallback if not running on Spaces
+
 class LLMService:
     def __init__(self):
         # Use Hugging Face Hub model name, not a local path
-        self.model_name = "Qwen/Qwen3-0.6B"
-        print("[LLMService] Initializing LLM Service...")
-        self.tokenizer = None
+        self.model_name = "Qwen/Qwen3-0.6B"  # Use the official repo name from the Hugging Face Hub
+        print("[LLMService] Loading Qwen tokenizer from Hugging Face Hub...")
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        print("[LLMService] Qwen tokenizer loaded.")
         self.model = None
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         print(f"[LLMService] Device set to: {self.device}")
-        
-        # Initialize tokenizer only, delay model loading
-        self._init_tokenizer()
-
-    def _init_tokenizer(self):
-        """Initialize tokenizer separately"""
-        try:
-            print("[LLMService] Loading Qwen tokenizer from Hugging Face Hub...")
-            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
-            if self.tokenizer.pad_token is None:
-                self.tokenizer.pad_token = self.tokenizer.eos_token
-            print("[LLMService] Qwen tokenizer loaded.")
-        except Exception as e:
-            print(f"[LLMService] Error loading tokenizer: {e}")
-            raise
 
     def load_model(self):
-        """Load model on-demand"""
+        """Load model without GPU decorator to avoid startup issues"""
         if self.model is None:
             print(f"[LLMService] Loading Qwen model from Hugging Face Hub on device: {self.device}...")
             try:
-                # Load with appropriate settings
-                load_kwargs = {
-                    "torch_dtype": torch.float16 if self.device == "cuda" else torch.float32,
-                    "trust_remote_code": True,
-                    "low_cpu_mem_usage": True
-                }
-                
-                if self.device == "cuda":
-                    load_kwargs["device_map"] = "auto"
-                
                 self.model = AutoModelForCausalLM.from_pretrained(
                     self.model_name,
-                    **load_kwargs
+                    torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
+                    device_map={"": self.device} if self.device == "cuda" else None,
+                    trust_remote_code=True  # Add this for some models
                 )
-                
                 if self.device == "cpu":
                     self.model = self.model.to(self.device)
-                
                 print(f"[LLMService] Qwen model loaded successfully on {self.device}.")
             except Exception as e:
                 print(f"[LLMService] Error loading model: {e}")
                 raise
 
+    @GPU  # Apply decorator to generation method instead
+    def _generate_with_gpu(self, inputs):
+        """GPU-accelerated generation method"""
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=300,
+                temperature=0.4,
+                top_p=0.8,
+                do_sample=True,
+                pad_token_id=self.tokenizer.eos_token_id  # Add this to avoid warnings
+            )
+        return outputs
+
     def _format_frequency_stats(self, stats: Dict) -> str:
         """Format frequency statistics into a readable string."""
         total_moves = sum(stats.get('move_distribution', {}).values())
@@ -68,6 +67,14 @@ class LLMService:
         for move, count in stats.get('move_distribution', {}).items():
             frequencies[move] = (count / total_moves) * 100
             
+        # Find the most common move
+        most_common = max(frequencies.items(), key=lambda x: x[1])
+        counter_move = {
+            "rock": "paper",
+            "paper": "scissors",
+            "scissors": "rock"
+        }[most_common[0]]
+            
         formatted = "AI Move Analysis:\n"
         formatted += f"Rock: {frequencies.get('rock', 0):.0f}%\n"
         formatted += f"Paper: {frequencies.get('paper', 0):.0f}%\n"
@@ -75,108 +82,89 @@ class LLMService:
         return formatted
 
     def _create_frequency_prompt(self, stats: Dict) -> str:
-        """Create a concise prompt for frequency-based analysis."""
+        """Create a concise prompt for frequency-based analysis with a beacon example."""
         base_prompt = (
-            "You are a helpful assistant that analyzes Rock-Paper-Scissors AI opponent statistics. "
-            "Based on the AI move frequency stats, recommend the best counter move. "
-            "Start your answer with '>>>' followed by explanation and 'Recommendation: [Move]'.\n\n"
+            "You are a helpful assistant that analyzes Rock-Paper-Scissors (RPS) AI opponent statistics to help players win. "
+            "Based on the provided AI move frequency stats, briefly recommend the best move for the player to counter the AI's most frequent move. "
+            "Your answer must start with '>>>', followed by a one-sentence explanation and then a new line with 'Recommendation: [Move]'.\n\n"
             "Example:\n"
-            ">>> The AI plays Scissors most (84%), so play Rock to counter it.\nRecommendation: Rock\n\n"
-            "AI Move Stats:\n"
+            ">>> The AI most often plays Scissors (84% of the time), so you should play Rock to counter it, since Rock beats Scissors.\nRecommendation: Rock\n\n"
+            "AI Move Frequency Stats:\n"
         )
         base_prompt += self._format_frequency_stats(stats)
         return base_prompt
 
     def _clean_llm_output(self, output: str) -> str:
-        """Clean and extract the relevant response."""
-        # Remove code blocks and markdown
+        # Remove code block markers and markdown
         output = re.sub(r'`{3,}', '', output)
-        output = re.sub(r'\*\*.*?\*\*', '', output)
-        
+        output = re.sub(r'\*\*Final Answer\*\*', '', output, flags=re.IGNORECASE)
+        # Find the beacon and extract the answer
         lines = [line.strip() for line in output.splitlines() if line.strip()]
-        
-        # Look for our beacon pattern
-        result_lines = []
         for i, line in enumerate(lines):
             if line.startswith('>>>'):
-                result_lines.append(line)
-                # Get the next line if it's a recommendation
-                if i + 1 < len(lines) and 'recommendation:' in lines[i + 1].lower():
-                    result_lines.append(lines[i + 1])
-                break
-        
-        if result_lines:
-            return '\n'.join(result_lines)
-        
-        # Fallback: look for any recommendation
+                result = [line]
+                # Optionally include the next line if it's a recommendation
+                if i + 1 < len(lines) and lines[i + 1].lower().startswith('recommendation:'):
+                    result.append(lines[i + 1])
+                return '\n'.join(result)
+        # Fallback: previous cleaning logic
+        seen = set()
+        cleaned_lines = []
         for line in lines:
-            if 'recommendation:' in line.lower():
+            if line not in seen:
+                cleaned_lines.append(line)
+                seen.add(line)
+        for line in cleaned_lines:
+            if line.lower().startswith('recommendation:'):
                 return line
-                
-        # Last resort fallback
-        return ">>> Try to counter the AI's most common move!\nRecommendation: Rock"
+        return '\n'.join(cleaned_lines)
 
     async def generate_response(self, prompt: str, stats: Optional[Dict] = None) -> str:
         print("[LLMService] generate_response called.")
-        
+        if stats:
+            print(f"[LLMService] Using stats for prompt: {stats}")
+            prompt = self._create_frequency_prompt(stats)
+        else:
+            print(f"[LLMService] Using raw prompt: {prompt}")
         try:
-            if stats:
-                print(f"[LLMService] Using stats for prompt: {stats}")
-                prompt = self._create_frequency_prompt(stats)
-            else:
-                print(f"[LLMService] Using raw prompt: {prompt}")
-            
             # Ensure model is loaded
             self.load_model()
             print("[LLMService] Model generation started...")
             
-            # Tokenize input
-            inputs = self.tokenizer([prompt], return_tensors="pt", padding=True, truncation=True)
-            
-            # Move to device if using GPU
-            if self.device == "cuda":
+            # Move inputs to the same device as model
+            inputs = self.tokenizer([prompt], return_tensors="pt")
+            if self.device == "cuda" and self.model is not None:
                 inputs = {k: v.to(self.device) for k, v in inputs.items()}
             
-            # Generate response
-            with torch.no_grad():
-                outputs = self.model.generate(
-                    **inputs,
-                    max_new_tokens=200,
-                    temperature=0.7,
-                    top_p=0.9,
-                    do_sample=True,
-                    pad_token_id=self.tokenizer.eos_token_id,
-                    eos_token_id=self.tokenizer.eos_token_id
-                )
+            # Use GPU-decorated method if available, otherwise regular generation
+            try:
+                outputs = self._generate_with_gpu(inputs)
+            except Exception as gpu_error:
+                print(f"[LLMService] GPU generation failed, falling back to CPU: {gpu_error}")
+                # Fallback to CPU generation
+                with torch.no_grad():
+                    outputs = self.model.generate(
+                        **inputs,
+                        max_new_tokens=300,
+                        temperature=0.4,
+                        top_p=0.8,
+                        do_sample=True,
+                        pad_token_id=self.tokenizer.eos_token_id
+                    )
             
-            # Decode response
             response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-            
-            # Extract only the new generated part (remove input prompt)
-            if prompt in response:
-                response = response.replace(prompt, "").strip()
-            
             print("[LLMService] Model generation finished.")
+            # Clean up the output before returning
             return self._clean_llm_output(response)
-            
         except Exception as e:
             print(f"[LLMService] Error during model generation: {e}")
-            # Return a reasonable fallback
-            if stats and 'move_distribution' in stats:
-                moves = stats['move_distribution']
-                most_common_move = max(moves.items(), key=lambda x: x[1])[0]
-                counter_moves = {"rock": "Paper", "paper": "Scissors", "scissors": "Rock"}
-                counter = counter_moves.get(most_common_move, "Rock")
-                return f">>> The AI often plays {most_common_move.title()}, so try {counter} to counter it!\nRecommendation: {counter}"
-            else:
-                return ">>> Try to observe the AI's patterns and counter accordingly!\nRecommendation: Rock"
+            # Return a fallback response instead of raising
+            return ">>> I'm having trouble analyzing the game right now, but generally try to counter the AI's most common move!\nRecommendation: Rock"
 
     async def close(self):
         """Clean up resources"""
         if self.model is not None:
             del self.model
-        if hasattr(self, 'tokenizer') and self.tokenizer is not None:
-            del self.tokenizer
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
         print("[LLMService] Resources cleaned up.")