Update app.py

app.py

@@ -8,20 +8,18 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 MODEL_ID = "Qwen/Qwen2-1.5B-Instruct"
 
 # --- Load Model and Tokenizer ---
-# Note: Model loading happens when the Space starts.
-# device_map="auto" will attempt to use the GPU when allocated by @spaces.GPU
 print(f"Loading model: {MODEL_ID}")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     torch_dtype="auto",
-    device_map="auto" # Keep this, it helps distribute within the allocated GPU(s)
+    device_map="auto"
 )
 print("Model loaded successfully.")
 
 
 # --- Generation Function (Returns response and token count) ---
-# This function will run on the GPU allocated via the decorator on process_input
+# No changes needed here
 def generate_response(messages, max_length=512, temperature=0.7, top_p=0.9):
     """Generate a response and return it along with the number of generated tokens."""
     num_generated_tokens = 0
@@ -31,8 +29,6 @@ def generate_response(messages, max_length=512, temperature=0.7, top_p=0.9):
             tokenize=False,
             add_generation_prompt=True
         )
-        # Ensure model_inputs are sent to the correct device the model is on
-        # device_map='auto' handles this, but explicitly checking model.device is safer
         device = model.device
         model_inputs = tokenizer([prompt_text], return_tensors="pt").to(device)
         input_ids_len = model_inputs.input_ids.shape[-1]
@@ -47,40 +43,48 @@ def generate_response(messages, max_length=512, temperature=0.7, top_p=0.9):
 
         print("Generating response...")
         with torch.no_grad():
-            # Generate response
             generated_ids = model.generate(model_inputs.input_ids, **generation_kwargs)
 
-        # 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(), num_generated_tokens
 
     except Exception as e:
         print(f"Error during generation: {e}")
-        # Ensure error message is returned correctly even if tokens couldn't be counted
         return f"An error occurred: {str(e)}", num_generated_tokens
 
-# --- Input Processing Function (Decorated for ZeroGPU) ---
-@spaces.GPU # Add the ZeroGPU decorator here
+# --- Input Processing Function (Adapts based on mode) ---
+@spaces.GPU # Keep ZeroGPU decorator
 def process_input(
+    analysis_mode, # New input: Mode selector
     player_stats,
-    ai_stats,
-    system_prompt,
-    user_query,
+    player_last_move,
+    markov_prediction_text,
+    system_prompt_freq, # Specific system prompt for frequency mode
+    system_prompt_markov, # Specific system prompt for markov mode
+    user_query, # User query might need slight adaptation based on mode
     max_length,
     temperature,
     top_p
 ):
-    """Process inputs, generate response, and return display info, response, time, and token count."""
-    print("GPU requested via decorator, starting processing...") # Add a log message
-    # Construct the user message content
-    user_content = f"Player Move Frequency Stats:\n{player_stats}\n\n"
-    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}"
+    """Process inputs based on selected analysis mode, generate response."""
+    print(f"GPU requested via decorator, starting processing in mode: {analysis_mode}")
+
+    # Select the appropriate system prompt and construct user content based on mode
+    if analysis_mode == "Frequency Only":
+        system_prompt = system_prompt_freq
+        user_content = f"Player Move Frequency Stats (Long-Term):\n{player_stats}\n\n"
+        user_content += f"User Query:\n{user_query}" # Query might need adjustment
+    elif analysis_mode == "Markov Prediction Only":
+        system_prompt = system_prompt_markov
+        user_content = f"Player's Last Move:\n{player_last_move}\n\n"
+        user_content += f"Predicted Next Move (Short-Term Markov Analysis):\n{markov_prediction_text}\n\n"
+        user_content += f"User Query:\n{user_query}" # Query might need adjustment
+    else:
+        # Default or error case
+        return "Invalid analysis mode selected.", "", "0 seconds", 0
 
     # Create the messages list
     messages = []
@@ -92,7 +96,7 @@ def process_input(
     start_time = time.time()
 
     # Generate response from the model
-    response, generated_tokens = generate_response(
+    response, generated_tokens = generate_response( # Capture token count
         messages,
         max_length=max_length,
         temperature=temperature,
@@ -101,71 +105,139 @@ def process_input(
 
     # --- Time Measurement End ---
     end_time = time.time()
-    duration = round(end_time - start_time, 2)
+    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}"
+    display_prompt = f"Selected Mode: {analysis_mode}\nSystem Prompt:\n{system_prompt}\n\n------\n\nUser Content:\n{user_content}"
 
-    print(f"Processing finished in {duration} seconds.") # Add a log message
+    print(f"Processing finished in {duration} seconds.")
     # Return all results including time and tokens
     return display_prompt, response, f"{duration} seconds", generated_tokens
 
-# --- Gradio Interface (No changes needed here) ---
+# --- System Prompts ---
+
+# Refined system prompt for Frequency Analysis
+DEFAULT_SYSTEM_PROMPT_FREQ = """You are an assistant that analyzes Rock-Paper-Scissors (RPS) player statistics. Your ONLY goal is to find the best single AI move to counter the player's MOST frequent move based on the provided frequency stats.
+
+Follow these steps EXACTLY. Do NOT deviate.
+
+Step 1: Identify Player's Most Frequent Move.
+   - Look ONLY at the 'Player Move Frequency Stats'.
+   - List the percentages: Rock (%), Paper (%), Scissors (%).
+   - State which move name has the highest percentage number.
+
+Step 2: Determine the Counter Move using RPS Rules.
+   - REMEMBER THE RULES: Paper beats Rock. Rock beats Scissors. Scissors beats Paper.
+   - Based *only* on the move identified in Step 1, state the single move name that beats it according to the rules. State the rule you used (e.g., "Paper beats Rock").
+
+Step 3: Explain the Counter Choice.
+   - Briefly state: "Playing [Counter Move from Step 2] is recommended because it directly beats the player's most frequent move, [Most Frequent Move from Step 1]."
+
+Step 4: State Final Recommendation.
+   - State *only* the recommended AI move name from Step 2. Example: "Recommendation: Paper"
 
-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.
+Base your analysis strictly on the provided frequencies and the stated RPS rules.
+"""
 
-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).
+# New system prompt for Markov Analysis
+DEFAULT_SYSTEM_PROMPT_MARKOV = """You are an assistant that analyzes Rock-Paper-Scissors (RPS) short-term player patterns. Your ONLY goal is to find the best single AI move to counter the player's PREDICTED next move, based on their LAST move.
 
-Base your analysis strictly on the provided frequencies and standard RPS rules."""
+Information Provided:
+1.  **Player's Last Move:** The actual move the player just made.
+2.  **Predicted Next Move (Short-Term Markov Analysis):** The player's statistically most likely *next* move based on their *last* move.
 
+Follow these steps EXACTLY:
+
+Step 1: Identify Predicted Player Move.
+    - Look at the 'Predicted Next Move (Short-Term Markov Analysis)' text.
+    - State the player's predicted next move (Rock, Paper, or Scissors). Note the probability if provided.
+
+Step 2: Determine Counter Move using RPS Rules.
+    - REMEMBER THE RULES: Paper beats Rock. Rock beats Scissors. Scissors beats Paper.
+    - Based *only* on the predicted move identified in Step 1, state the single AI move name that beats it. State the rule used (e.g., "Rock beats Scissors").
+
+Step 3: Explain the Counter Choice.
+    - Briefly state: "Playing [Counter Move from Step 2] is recommended because it directly beats the player's predicted next move, [Predicted Move from Step 1]."
+
+Step 4: State Final Recommendation.
+    - State *only* the recommended AI move name from Step 2. Example: "Recommendation: Rock"
+
+Base your analysis strictly on the provided prediction and the standard RPS rules.
+"""
+
+# --- Default Input Values ---
 DEFAULT_PLAYER_STATS = "Rock: 40%\nPaper: 30%\nScissors: 30%"
-DEFAULT_AI_STATS = ""
-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."
+DEFAULT_PLAYER_LAST_MOVE = "Rock"
+DEFAULT_MARKOV_PREDICTION = "Based on the last move (Rock), the player's most likely next move is Paper (60% probability)."
+# Default query might need to be generic or adapted based on mode
+DEFAULT_USER_QUERY = "Based on the provided information for the selected analysis mode, what single move should the AI make next? Explain your reasoning step-by-step as instructed."
 
+# --- Gradio Interface ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown(f"# {MODEL_ID} - RPS Frequency Analysis Tester")
-    gr.Markdown("Test model advice based on Player/AI move frequencies. Includes Generation Time and Token Count.")
+    gr.Markdown(f"# {MODEL_ID} - RPS Strategy Tester")
+    gr.Markdown("Test model advice using either Frequency Stats OR Short-Term (Markov) Predictions.")
+
+    # Mode Selector
+    analysis_mode_selector = gr.Radio(
+        label="Select Analysis Mode",
+        choices=["Frequency Only", "Markov Prediction Only"],
+        value="Frequency Only" # Default mode
+    )
 
+    # Input Sections (conditionally visible)
+    with gr.Group(visible=True) as frequency_inputs: # Visible by default
+        gr.Markdown("### Frequency Analysis Inputs")
+        player_stats_input = gr.Textbox(
+            label="Player Move Frequency Stats (Long-Term)", value=DEFAULT_PLAYER_STATS, lines=4,
+            info="Overall player move distribution."
+        )
+        # Hidden system prompt for frequency mode (can be edited if needed)
+        system_prompt_freq_input = gr.Textbox(
+            label="System Prompt (Frequency Mode - Edit if needed)", value=DEFAULT_SYSTEM_PROMPT_FREQ, lines=15, visible=False # Hidden by default, but can be shown for advanced editing
+        )
+
+    with gr.Group(visible=False) as markov_inputs: # Hidden by default
+        gr.Markdown("### Markov Prediction Analysis Inputs")
+        player_last_move_input = gr.Dropdown( # Dropdown is good for defined choices
+            label="Player's Last Move", choices=["Rock", "Paper", "Scissors"], value=DEFAULT_PLAYER_LAST_MOVE,
+            info="The player's most recent actual move."
+        )
+        markov_prediction_input = gr.Textbox(
+            label="Predicted Next Move (Short-Term Markov Analysis)", value=DEFAULT_MARKOV_PREDICTION, lines=3,
+            info="Provide the pre-calculated prediction based on the last move (e.g., 'Player likely plays Paper (60%)')."
+        )
+         # Hidden system prompt for markov mode (can be edited if needed)
+        system_prompt_markov_input = gr.Textbox(
+            label="System Prompt (Markov Mode - Edit if needed)", value=DEFAULT_SYSTEM_PROMPT_MARKOV, lines=15, visible=False # Hidden by default
+        )
+
+    # General Inputs / Parameters / Outputs
     with gr.Row():
-        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 player's move frequencies (e.g., Rock: 50%, Paper: 30%, Scissors: 20%)."
-            )
-            ai_stats_input = gr.Textbox(
-                label="AI Move Frequency Stats (Optional)", value=DEFAULT_AI_STATS, lines=4,
-                info="Optionally, enter AI's own move frequencies."
-            )
-            user_query_input = gr.Textbox(
-                label="Your Query / Instruction", value=DEFAULT_USER_QUERY, lines=3,
-                info="Ask the specific question based on the stats."
-            )
-            system_prompt_input = gr.Textbox(
-                label="System Prompt", value=DEFAULT_SYSTEM_PROMPT,
-                lines=12
-            )
-
-        with gr.Column(scale=1): # Params/Output column
-            gr.Markdown("## Generation Parameters")
+         with gr.Column(scale=2):
+             user_query_input = gr.Textbox(
+                 label="Your Query / Instruction", value=DEFAULT_USER_QUERY, lines=3,
+                 info="Ask the specific question based on the selected mode's analysis."
+             )
+         with gr.Column(scale=1):
+            gr.Markdown("#### Generation Parameters")
             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")
             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")
 
-            gr.Markdown("## Performance Metrics")
-            time_output = gr.Textbox(label="Generation Time", interactive=False)
-            tokens_output = gr.Number(label="Generated Tokens", interactive=False)
 
-            gr.Markdown("""
-            ## Testing Tips
-            - Focus on player stats for optimal counter strategy.
-            - Use the refined **System Prompt** for better reasoning guidance.
-            - Lower **Temperature** encourages more direct, statistical answers.
+    submit_btn = gr.Button("Generate Response", variant="primary")
+
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("#### Performance Metrics")
+            time_output = gr.Textbox(label="Generation Time", interactive=False)
+            tokens_output = gr.Number(label="Generated Tokens", interactive=False) # Use Number for token count
+        with gr.Column():
+             gr.Markdown("""
+            #### Testing Tips
+            - Select the desired **Analysis Mode**.
+            - Fill in the inputs for the **selected mode only**.
+            - Use low **Temperature** for factual analysis.
             """)
 
     with gr.Row():
@@ -176,21 +248,56 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             label="Model Response", lines=20, show_copy_button=True
         )
 
+    # --- Event Handlers ---
+
+    # Function to update UI visibility based on mode selection
+    def update_ui_visibility(mode):
+        if mode == "Frequency Only":
+            return {
+                frequency_inputs: gr.update(visible=True),
+                markov_inputs: gr.update(visible=False)
+            }
+        elif mode == "Markov Prediction Only":
+            return {
+                frequency_inputs: gr.update(visible=False),
+                markov_inputs: gr.update(visible=True)
+            }
+        else: # Default case
+             return {
+                frequency_inputs: gr.update(visible=True),
+                markov_inputs: gr.update(visible=False)
+            }
+
+    # Link the radio button change to the UI update function
+    analysis_mode_selector.change(
+        fn=update_ui_visibility,
+        inputs=analysis_mode_selector,
+        outputs=[frequency_inputs, markov_inputs] # Components to update
+    )
+
+    # Handle button click - Pass all inputs, function will select based on mode
     submit_btn.click(
         process_input,
         inputs=[
-            player_stats_input, ai_stats_input, system_prompt_input,
-            user_query_input, max_length_slider, temperature_slider, top_p_slider
+            analysis_mode_selector, # Mode selector first
+            player_stats_input,
+            player_last_move_input,
+            markov_prediction_input,
+            system_prompt_freq_input, # Pass both system prompts
+            system_prompt_markov_input,
+            user_query_input,
+            max_length_slider,
+            temperature_slider,
+            top_p_slider
         ],
         outputs=[
             final_prompt_display, response_display,
             time_output, tokens_output
         ],
-        api_name="generate_rps_frequency_analysis_v2"
+        api_name="generate_rps_selectable_analysis" # Updated api_name
     )
 
 # --- Launch the demo ---
 if __name__ == "__main__":
-    # Share=True is needed for ZeroGPU to work correctly if running locally for testing
-    # but usually not needed when deployed on HF Spaces platform.
+    # Share=True might be needed for ZeroGPU if running locally for testing
     demo.launch()