Update app.py

app.py

@@ -2,154 +2,152 @@ import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Load model and tokenizer
-model_name = "Qwen/Qwen2-0.5B"
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+# Load the Qwen2 0.5B model
+model_id = "Qwen/Qwen2-0.5B"
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    device_map="auto",
+    model_id,
     torch_dtype=torch.float16,
+    device_map="auto",
     trust_remote_code=True
 )
 
-def generate_response(prompt, max_length=300, temperature=0.7):
+def generate_response(prompt, max_length=512, temperature=0.7, top_p=0.9):
+    """Generate a response from the Qwen2 model based on the input prompt."""
+    
+    # Tokenize the input prompt
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     
+    # Generate response
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
             max_new_tokens=max_length,
             do_sample=True,
             temperature=temperature,
-            top_p=0.9,
+            top_p=top_p,
         )
     
+    # Decode the response
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    # Remove the prompt from the response
-    if response.startswith(prompt):
+    
+    # Extract only the model's response (remove the input prompt)
+    if prompt in response:
         response = response[len(prompt):]
     
     return response.strip()
 
-# 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}
-
-Based on this information, analyze the game and recommend a next move.
-""",
+def process_input(
+    raw_prompt, 
+    game_stats_template,
+    template_type,
+    max_length, 
+    temperature, 
+    top_p
+):
+    """Process the input and template to create the final prompt for the model."""
     
-    "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}%
-
-What should be the AI's next move?
-""",
+    final_prompt = ""
     
-    "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:
-"""
-}
-
-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 {}
-
-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 = []
+    if template_type == "Raw Prompt Only":
+        final_prompt = raw_prompt
+    elif template_type == "Template + Prompt":
+        final_prompt = f"{game_stats_template}\n\n{raw_prompt}"
+    elif template_type == "Custom Format":
+        final_prompt = f"{game_stats_template}\n\nBased on the game statistics above, {raw_prompt}"
     
-    for key, value in sample_data.items():
-        inputs.append(gr.Textbox(value=value, label=key))
-    
-    return inputs
-
-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)
+    # Generate response from the model
+    response = generate_response(
+        final_prompt,
+        max_length=max_length,
+        temperature=temperature,
+        top_p=top_p
+    )
     
-    return prompt, response
+    return final_prompt, response
 
-# Define the interface
+# Create the Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Qwen2 0.5B Testing for Rock-Paper-Scissors Game Analysis")
+    gr.Markdown("# Qwen2 0.5B Game Analysis Tester")
+    gr.Markdown("Use this interface to test how the Qwen2 0.5B model responds to different prompts about your game statistics.")
     
     with gr.Row():
         with gr.Column():
-            template_dropdown = gr.Dropdown(
-                choices=list(test_templates.keys()),
-                value="Basic Game History",
-                label="Select Template"
+            template_type = gr.Radio(
+                ["Raw Prompt Only", "Template + Prompt", "Custom Format"],
+                label="Prompt Template Type",
+                value="Template + Prompt"
             )
             
-            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()]
+            game_stats_template = gr.Textbox(
+                label="Game Statistics Template",
+                placeholder="Enter your game statistics here (scores, round history, etc.)",
+                lines=10
+            )
+            
+            raw_prompt = gr.Textbox(
+                label="Prompt",
+                placeholder="What do you want the model to analyze or respond to?",
+                lines=3
+            )
             
-            for field in input_fields:
-                input_container.append(field)
+            with gr.Row():
+                max_length = gr.Slider(
+                    minimum=50, 
+                    maximum=1024, 
+                    value=256, 
+                    step=1, 
+                    label="Max Response Length"
+                )
+                temperature = gr.Slider(
+                    minimum=0.1, 
+                    maximum=1.5, 
+                    value=0.7, 
+                    step=0.1, 
+                    label="Temperature"
+                )
+                top_p = gr.Slider(
+                    minimum=0.1, 
+                    maximum=1.0, 
+                    value=0.9, 
+                    step=0.1, 
+                    label="Top P"
+                )
                 
-            test_button = gr.Button("Test Model")
+            submit_btn = gr.Button("Generate Response")
         
         with gr.Column():
-            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()]
+            final_prompt_display = gr.Textbox(
+                label="Final Prompt Sent to Model", 
+                lines=10
+            )
+            response_display = gr.Textbox(
+                label="Model Response", 
+                lines=15
+            )
     
-    template_dropdown.change(
-        fn=update_inputs,
-        inputs=template_dropdown,
-        outputs=input_container
+    submit_btn.click(
+        process_input,
+        inputs=[
+            raw_prompt,
+            game_stats_template,
+            template_type,
+            max_length,
+            temperature,
+            top_p
+        ],
+        outputs=[final_prompt_display, response_display]
     )
     
-    test_button.click(
-        fn=test_model,
-        inputs=[template_dropdown] + input_fields,
-        outputs=[prompt_output, response_output]
-    )
+    gr.Markdown("""
+    ## Tips for Testing
+    
+    1. Start with simple prompts to gauge the model's basic understanding
+    2. Gradually increase complexity to find the model's limitations
+    3. Try different prompt formats to see which works best
+    4. Experiment with temperature and top_p to find optimal settings
+    5. Document which prompts work well as candidates for fine-tuning
+    """)
 
+# Launch the demo
 demo.launch()