Update app.py

app.py

@@ -1,132 +1,155 @@
 import gradio as gr
-import json
-import collections
+import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Define model name - use a very small model
-MODEL_NAME = "Qwen/Qwen2-0.5B-Instruct"  # 0.5B parameter instruction model
-
-print(f"Loading model {MODEL_NAME}...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+# Load model and tokenizer
+model_name = "Qwen/Qwen2-0.5B"
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    low_cpu_mem_usage=True,  # CPU-friendly settings
-    device_map="cpu"  # Force CPU usage
+    model_name,
+    device_map="auto",
+    torch_dtype=torch.float16,
+    trust_remote_code=True
 )
-print("Model loaded successfully!")
-
-def format_rps_game_prompt(game_data):
-    """Format Rock-Paper-Scissors game data into a simple prompt for the LLM"""
-    try:
-        # Parse the JSON game data
-        if isinstance(game_data, str):
-            game_data = json.loads(game_data)
-        
-        # Extract key game information
-        player_history = game_data.get("player_history", [])
-        opponent_history = game_data.get("opponent_history", [])
-        rounds_played = len(player_history)
-        player_score = game_data.get("player_score", 0)
-        opponent_score = game_data.get("opponent_score", 0)
-        draws = game_data.get("draws", 0)
-        
-        # Create a simple prompt with just the game state
-        prompt = f"""You are a Rock-Paper-Scissors strategy advisor. 
 
-Game State:
-- Rounds played: {rounds_played}
-- Player score: {player_score}
-- Opponent score: {opponent_score}
-- Draws: {draws}
-- Player's moves (oldest to newest): {', '.join(player_history)}
-- Opponent's moves (oldest to newest): {', '.join(opponent_history)}
+def generate_response(prompt, max_length=300, temperature=0.7):
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=max_length,
+            do_sample=True,
+            temperature=temperature,
+            top_p=0.9,
+        )
+    
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # Remove the prompt from the response
+    if response.startswith(prompt):
+        response = response[len(prompt):]
+    
+    return response.strip()
 
-Analyze the game history and provide advice on what move (Rock, Paper, or Scissors) the player should make next.
+# Create different test templates
+test_templates = {
+    "Basic Game History": """
+Game history: {game_history}
+Player score: {player_score}
+AI score: {ai_score}
+Last move: {last_move}
 
-First, explain your thought process and reasoning in detail:
-1. Look for patterns in the opponent's moves
-2. Consider if the opponent seems to be using any strategy
-3. Evaluate if the player should try to counter the opponent's most likely next move
-4. Consider any psychological factors that might be relevant
+Based on this information, analyze the game and recommend a next move.
+""",
+    
+    "With Pre-calculated Statistics": """
+Game history: {game_history}
+Player's move frequencies: Rock ({rock_freq}%), Paper ({paper_freq}%), Scissors ({scissors_freq}%)
+Player's patterns: 
+- After playing Rock, chooses Paper: {rock_to_paper}% 
+- After playing Paper, chooses Scissors: {paper_to_scissors}%
+- After playing Scissors, chooses Rock: {scissors_to_rock}%
 
-After your explanation, end with a clear recommendation: "Recommendation: [Rock/Paper/Scissors]"
+What should be the AI's next move?
+""",
+    
+    "Simplified Decision": """
+Recent moves: {recent_moves}
+Based on this pattern, the player is likely to play {likely_next} next.
+To counter {likely_next}, the AI should play:
 """
-        return prompt
-    except Exception as e:
-        return f"Error formatting prompt: {str(e)}\n\nPlease provide game data in a valid JSON format."
+}
 
-def generate_advice(game_data):
-    """Generate advice based on game data using the LLM"""
-    try:
-        # Format the prompt
-        prompt = format_rps_game_prompt(game_data)
-        
-        # Generate response from LLM (with CPU-only settings)
-        inputs = tokenizer(prompt, return_tensors="pt")
-        
-        # Set max_length to avoid excessive generation
-        outputs = model.generate(
-            inputs["input_ids"],
-            max_new_tokens=200,  # Allow more tokens for detailed reasoning
-            do_sample=True, 
-            temperature=0.7, 
-            top_p=0.9
-        )
-        full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        
-        # Extract just the model's response (remove the prompt)
-        if full_response.startswith(prompt):
-            response = full_response[len(prompt):].strip()
-        else:
-            response = full_response
-            
-        # If model response is too short, add fallback advice
-        if len(response) < 30:
-            # Simple fallback based on opponent's last move
-            if len(opponent_history) > 0:
-                last_move = opponent_history[-1]
-                if last_move == "Rock":
-                    suggestion = "Paper"
-                elif last_move == "Paper":
-                    suggestion = "Scissors"
-                else:
-                    suggestion = "Rock"
-                    
-                response += f"\n\nBased on the opponent's last move ({last_move}), a reasonable counter would be:\nRecommendation: {suggestion}"
-            else:
-                response += "\n\nWith no game history to analyze, each move has equal probability of success.\nRecommendation: Paper"
-        
-        return response
-    except Exception as e:
-        return f"Error generating advice: {str(e)}"
+def create_sample_data(template_key):
+    """Create sample data for the selected template"""
+    if template_key == "Basic Game History":
+        return {
+            "game_history": "R,P,S,R,P,S,S,R,P,R",
+            "player_score": "5",
+            "ai_score": "3",
+            "last_move": "P"
+        }
+    elif template_key == "With Pre-calculated Statistics":
+        return {
+            "game_history": "R,P,S,R,P,S,S,R,P,R",
+            "rock_freq": "40",
+            "paper_freq": "30",
+            "scissors_freq": "30",
+            "rock_to_paper": "75",
+            "paper_to_scissors": "67",
+            "scissors_to_rock": "50"
+        }
+    elif template_key == "Simplified Decision":
+        return {
+            "recent_moves": "R,P,S,R,P",
+            "likely_next": "S"
+        }
+    return {}
 
-# Sample game data for the example
-sample_game_data = {
-    "player_history": ["Rock", "Paper", "Scissors", "Rock", "Paper"],
-    "opponent_history": ["Scissors", "Rock", "Paper", "Scissors", "Rock"],
-    "player_score": 3,
-    "opponent_score": 2,
-    "draws": 0
-}
+def format_prompt(template_key, **kwargs):
+    """Format the selected template with provided values"""
+    template = test_templates[template_key]
+    return template.format(**kwargs)
+
+def update_template_inputs(template_name):
+    """Update the input fields based on the selected template"""
+    sample_data = create_sample_data(template_name)
+    inputs = []
+    
+    for key, value in sample_data.items():
+        inputs.append(gr.Textbox(value=value, label=key))
+    
+    return inputs
 
-# Create Gradio interface
-with gr.Blocks(title="Rock-Paper-Scissors AI Assistant") as demo:
-    gr.Markdown("# Rock-Paper-Scissors AI Assistant")
-    gr.Markdown("Enter your game data to get advice on your next move.")
+def test_model(template_name, *args):
+    """Test the model with the provided template and inputs"""
+    sample_data = create_sample_data(template_name)
+    data = dict(zip(sample_data.keys(), args))
+    
+    prompt = format_prompt(template_name, **data)
+    response = generate_response(prompt)
+    
+    return prompt, response
+
+# Define the interface
+with gr.Blocks() as demo:
+    gr.Markdown("# Qwen2 0.5B Testing for Rock-Paper-Scissors Game Analysis")
     
     with gr.Row():
         with gr.Column():
-            game_data_input = gr.Textbox(
-                label="Game State (JSON)",
-                placeholder=json.dumps(sample_game_data, indent=2),
-                lines=10
+            template_dropdown = gr.Dropdown(
+                choices=list(test_templates.keys()),
+                value="Basic Game History",
+                label="Select Template"
             )
-            submit_btn = gr.Button("Get Advice")
+            
+            input_container = gr.Column()
+            sample_data = create_sample_data("Basic Game History")
+            input_fields = [gr.Textbox(value=v, label=k) for k, v in sample_data.items()]
+            
+            for field in input_fields:
+                input_container.append(field)
+                
+            test_button = gr.Button("Test Model")
         
         with gr.Column():
-            output = gr.Textbox(label="AI Advice", lines=10)
+            prompt_output = gr.Textbox(label="Formatted Prompt")
+            response_output = gr.Textbox(label="Model Response")
+    
+    def update_inputs(template_name):
+        sample_data = create_sample_data(template_name)
+        return [gr.Textbox(value=v, label=k) for k, v in sample_data.items()]
+    
+    template_dropdown.change(
+        fn=update_inputs,
+        inputs=template_dropdown,
+        outputs=input_container
+    )
     
-    submit_btn.click(generate_advice, inputs=[game_data_input], outputs=[output])
+    test_button.click(
+        fn=test_model,
+        inputs=[template_dropdown] + input_fields,
+        outputs=[prompt_output, response_output]
+    )
 
-# Launch the app
 demo.launch()