Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import torch
+import time # Import time module
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # --- Configuration ---
@@ -16,9 +17,10 @@ model = AutoModelForCausalLM.from_pretrained(
 print("Model loaded successfully.")
 
 
-# --- Generation Function (Using Chat Template - No changes needed here) ---
+# --- Generation Function (Updated to return token count) ---
 def generate_response(messages, max_length=512, temperature=0.7, top_p=0.9):
-    """Generate a response from the Qwen2 model using chat template."""
+    """Generate a response and return it along with the number of generated tokens."""
+    num_generated_tokens = 0
     try:
         prompt_text = tokenizer.apply_chat_template(
             messages,
@@ -26,6 +28,8 @@ def generate_response(messages, max_length=512, temperature=0.7, top_p=0.9):
             add_generation_prompt=True
         )
         model_inputs = tokenizer([prompt_text], return_tensors="pt").to(model.device)
+        input_ids_len = model_inputs.input_ids.shape[-1] # Length of input tokens
+
         generation_kwargs = {
             "max_new_tokens": max_length,
             "temperature": temperature,
@@ -33,155 +37,151 @@ def generate_response(messages, max_length=512, temperature=0.7, top_p=0.9):
             "do_sample": True,
             "pad_token_id": tokenizer.eos_token_id,
         }
+
         print("Generating response...")
         with torch.no_grad():
+            # Generate response - Ensure output_scores or similar isn't needed if just counting
             generated_ids = model.generate(model_inputs.input_ids, **generation_kwargs)
-        input_ids_len = model_inputs.input_ids.shape[-1]
+
+        # Calculate generated tokens
         output_ids = generated_ids[0, input_ids_len:]
+        num_generated_tokens = len(output_ids)
+
         response = tokenizer.decode(output_ids, skip_special_tokens=True)
         print("Generation complete.")
-        return response.strip()
+        return response.strip(), num_generated_tokens # Return response and token count
 
     except Exception as e:
         print(f"Error during generation: {e}")
-        return f"An error occurred: {str(e)}"
+        return f"An error occurred: {str(e)}", num_generated_tokens # Return error and token count
 
-# --- Input Processing Function (Simplified for Frequency Stats) ---
+# --- Input Processing Function (Updated for Time/Token outputs) ---
 def process_input(
     player_stats,
     ai_stats,
     system_prompt,
-    user_query, # Changed from prompt template/custom
+    user_query,
     max_length,
     temperature,
     top_p
 ):
-    """Process frequency stats and user query for the model."""
+    """Process inputs, generate response, and return display info, response, time, and token count."""
 
-    # Construct the user message content using the provided stats and query
+    # Construct the user message content
     user_content = f"Player Move Frequency Stats:\n{player_stats}\n\n"
-    if ai_stats and ai_stats.strip(): # Include AI stats if provided
+    if ai_stats and ai_stats.strip():
          user_content += f"AI Move Frequency Stats:\n{ai_stats}\n\n"
     user_content += f"User Query:\n{user_query}"
 
-    # Create the messages list for the chat template
+    # Create the messages list
     messages = []
-    if system_prompt and system_prompt.strip(): # Add system prompt if provided
+    if system_prompt and system_prompt.strip():
         messages.append({"role": "system", "content": system_prompt})
     messages.append({"role": "user", "content": user_content})
 
+    # --- Time Measurement Start ---
+    start_time = time.time()
+
     # Generate response from the model
-    response = generate_response(
+    response, generated_tokens = generate_response( # Capture token count
         messages,
         max_length=max_length,
         temperature=temperature,
         top_p=top_p
     )
 
-    # For display purposes, show the constructed input
+    # --- Time Measurement End ---
+    end_time = time.time()
+    duration = round(end_time - start_time, 2) # Calculate duration
+
+    # For display purposes
     display_prompt = f"System Prompt (if used):\n{system_prompt}\n\n------\n\nUser Content:\n{user_content}"
 
-    return display_prompt, response
+    # Return all results including time and tokens
+    return display_prompt, response, f"{duration} seconds", generated_tokens
+
+# --- Gradio Interface (Added Time/Token displays, refined System Prompt) ---
 
-# --- Gradio Interface (Simplified for Frequency Stats) ---
+# Refined default system prompt for better reasoning
+DEFAULT_SYSTEM_PROMPT = """You are an expert Rock-Paper-Scissors (RPS) strategist focusing on statistical analysis.
+Your task is to recommend the optimal AI move based *only* on the provided move frequency statistics for the player.
 
-# Define a default system prompt suitable for frequency analysis
-DEFAULT_SYSTEM_PROMPT = """You are an expert Rock-Paper-Scissors (RPS) strategist.
-Analyze the provided frequency statistics for the player's (and potentially AI's) past moves.
-Based *only* on these statistics, determine the statistically optimal counter-strategy or recommendation for the AI's next move.
-Explain your reasoning clearly based on the probabilities implied by the frequencies and the rules of RPS (Rock beats Scissors, Scissors beats Paper, Paper beats Rock).
-Provide a clear recommendation (Rock, Paper, or Scissors) and justify it using expected outcomes or probabilities."""
+Follow these steps:
+1.  **Identify Player's Most Frequent Move:** Note the move (Rock, Paper, or Scissors) the player uses most often according to the stats.
+2.  **Determine Best Counter:** Identify the RPS move that directly beats the player's most frequent move (Rock beats Scissors, Scissors beats Paper, Paper beats Rock).
+3.  **Justify Recommendation:** Explain *why* this counter-move is statistically optimal. You can mention the expected outcome. For example: 'Playing Paper counters the player's most frequent move, Rock (40% frequency). This offers the highest probability of winning against the player's likely action.' Avoid irrelevant justifications based on the AI's own move frequencies.
+4.  **State Recommendation:** Clearly state the recommended move (Rock, Paper, or Scissors).
 
-# Default example stats
+Base your analysis strictly on the provided frequencies and standard RPS rules."""
+
+# Default example stats and query
 DEFAULT_PLAYER_STATS = "Rock: 40%\nPaper: 30%\nScissors: 30%"
-DEFAULT_AI_STATS = "Rock: 33%\nPaper: 34%\nScissors: 33%" # Example AI stats
-DEFAULT_USER_QUERY = "Based on the player's move frequencies, what move should the AI make next to maximize its statistical chance of winning? Explain your reasoning."
+DEFAULT_AI_STATS = "" # Keep AI stats optional and clear by default
+DEFAULT_USER_QUERY = "Based *only* on the player's move frequencies, what single move should the AI make next to maximize its statistical chance of winning? Explain your reasoning clearly step-by-step as instructed."
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(f"# {MODEL_ID} - RPS Frequency Analysis Tester")
-    gr.Markdown("Test how the model provides strategic advice based *only* on Player and AI move frequency statistics.")
+    gr.Markdown("Test model advice based on Player/AI move frequencies. Includes Generation Time and Token Count.")
 
     with gr.Row():
-        with gr.Column(scale=2): # Make input column wider
-            # Input for Player Stats
+        with gr.Column(scale=2): # Input column
             player_stats_input = gr.Textbox(
-                label="Player Move Frequency Stats",
-                value=DEFAULT_PLAYER_STATS,
-                lines=4,
-                info="Enter the observed frequencies of the player's moves."
+                label="Player Move Frequency Stats", value=DEFAULT_PLAYER_STATS, lines=4,
+                info="Enter player's move frequencies (e.g., Rock: 50%, Paper: 30%, Scissors: 20%)."
             )
-             # Input for AI Stats (Optional)
             ai_stats_input = gr.Textbox(
-                label="AI Move Frequency Stats (Optional)",
-                value=DEFAULT_AI_STATS,
-                lines=4,
-                info="Optionally, enter the AI's own move frequencies if relevant."
+                label="AI Move Frequency Stats (Optional)", value=DEFAULT_AI_STATS, lines=4,
+                info="Optionally, enter AI's own move frequencies."
             )
-            # Input for User Query
             user_query_input = gr.Textbox(
-                label="Your Query / Instruction",
-                value=DEFAULT_USER_QUERY,
-                lines=3,
-                info="Ask the specific question based on the frequency stats."
+                label="Your Query / Instruction", value=DEFAULT_USER_QUERY, lines=3,
+                info="Ask the specific question based on the stats."
             )
-            # System prompt (optional)
             system_prompt_input = gr.Textbox(
-                label="System Prompt (Optional)",
-                value=DEFAULT_SYSTEM_PROMPT,
-                placeholder="Define the AI's role and task based on frequency stats...",
-                lines=10 # Reduced lines needed
+                label="System Prompt", value=DEFAULT_SYSTEM_PROMPT, # Set default value
+                lines=12 # Adjusted lines
             )
 
-        with gr.Column(scale=1): # Make params/output column narrower
-             # Generation parameters
+        with gr.Column(scale=1): # Params/Output column
             gr.Markdown("## Generation Parameters")
-            max_length_slider = gr.Slider(
-                minimum=50, maximum=1024, value=350, step=16, label="Max New Tokens" # Reduced default length needed
-            )
-            temperature_slider = gr.Slider(
-                minimum=0.1, maximum=1.5, value=0.5, step=0.05, label="Temperature" # Defaulting lower for stats analysis
-            )
-            top_p_slider = gr.Slider(
-                minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top P"
-            )
-            # Generate button
+            max_length_slider = gr.Slider(minimum=50, maximum=1024, value=300, step=16, label="Max New Tokens")
+            temperature_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.4, step=0.05, label="Temperature") # Lowered default further
+            top_p_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top P")
             submit_btn = gr.Button("Generate Response", variant="primary")
 
-            # Tips for using the interface
+            gr.Markdown("## Performance Metrics")
+            # Outputs for Time and Tokens
+            time_output = gr.Textbox(label="Generation Time", interactive=False)
+            tokens_output = gr.Number(label="Generated Tokens", interactive=False) # Use Number for token count
+
             gr.Markdown("""
             ## Testing Tips
-            - Input player move frequencies directly. AI stats are optional.
-            - Refine the **User Query** to guide the model's task.
-            - Adjust the **System Prompt** for role/task definition.
-            - Use lower **Temperature** for more deterministic, calculation-like responses based on stats.
+            - Focus on player stats for optimal counter strategy.
+            - Use the refined **System Prompt** for better reasoning guidance.
+            - Lower **Temperature** encourages more direct, statistical answers.
             """)
 
     with gr.Row():
-         with gr.Column():
-            # Display final prompt and model response
-            final_prompt_display = gr.Textbox(
-                label="Formatted Input Sent to Model (via Chat Template)", lines=15
-            )
-         with gr.Column():
-            response_display = gr.Textbox(
-                label="Model Response", lines=15, show_copy_button=True
-            )
-
+        # Display final prompt and model response (side-by-side)
+        final_prompt_display = gr.Textbox(
+            label="Formatted Input Sent to Model (via Chat Template)", lines=20 # Increased lines
+        )
+        response_display = gr.Textbox(
+            label="Model Response", lines=20, show_copy_button=True # Increased lines
+        )
 
-    # Handle button click - Updated inputs list
+    # Handle button click - Updated inputs and outputs list
     submit_btn.click(
         process_input,
         inputs=[
-            player_stats_input,
-            ai_stats_input,
-            system_prompt_input,
-            user_query_input, # New input
-            max_length_slider,
-            temperature_slider,
-            top_p_slider
+            player_stats_input, ai_stats_input, system_prompt_input,
+            user_query_input, max_length_slider, temperature_slider, top_p_slider
+        ],
+        outputs=[
+            final_prompt_display, response_display,
+            time_output, tokens_output # Added new outputs
         ],
-        outputs=[final_prompt_display, response_display],
-        api_name="generate_rps_frequency_analysis" # Updated api_name
+        api_name="generate_rps_frequency_analysis_v2" # Updated api_name
     )
 
 # --- Launch the demo ---