Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| import time | |
| import spaces # Import the spaces library | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| # --- Configuration --- | |
| MODEL_ID = "Qwen/Qwen2.5-Math-1.5B" # Replace with actual ID if found | |
| # --- Load Model and Tokenizer --- | |
| print(f"Loading model: {MODEL_ID}") | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| MODEL_ID, | |
| torch_dtype="auto", | |
| device_map="auto" | |
| ) | |
| print("Model loaded successfully.") | |
| # --- Generation Function (Returns response and token count) --- | |
| # 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 | |
| try: | |
| prompt_text = tokenizer.apply_chat_template( | |
| messages, | |
| tokenize=False, | |
| add_generation_prompt=True | |
| ) | |
| device = model.device | |
| model_inputs = tokenizer([prompt_text], return_tensors="pt").to(device) | |
| input_ids_len = model_inputs.input_ids.shape[-1] | |
| generation_kwargs = { | |
| "max_new_tokens": max_length, | |
| "temperature": temperature, | |
| "top_p": top_p, | |
| "do_sample": True, | |
| "pad_token_id": tokenizer.eos_token_id, | |
| } | |
| print("Generating response...") | |
| with torch.no_grad(): | |
| generated_ids = model.generate(model_inputs.input_ids, **generation_kwargs) | |
| 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}") | |
| return f"An error occurred: {str(e)}", num_generated_tokens | |
| # --- Input Processing Function (Takes single system prompt) --- | |
| # Keep ZeroGPU decorator | |
| def process_input( | |
| analysis_mode, # Mode selector | |
| player_stats, | |
| player_last_move, | |
| markov_prediction_text, | |
| system_prompt, # Single system prompt from UI | |
| user_query, | |
| max_length, | |
| temperature, | |
| top_p | |
| ): | |
| """Process inputs based on selected analysis mode using the provided system prompt.""" | |
| print(f"GPU requested via decorator, starting processing in mode: {analysis_mode}") | |
| # Construct user content based on mode | |
| if analysis_mode == "Frequency Only": | |
| user_content = f"Player Move Frequency Stats (Long-Term):\n{player_stats}\n\n" | |
| user_content += f"User Query:\n{user_query}" | |
| elif analysis_mode == "Markov Prediction Only": | |
| 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}" | |
| else: | |
| return "Invalid analysis mode selected.", "", "0 seconds", 0 | |
| # Create the messages list using the system_prompt from the UI | |
| messages = [] | |
| 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, generated_tokens = generate_response( | |
| messages, | |
| max_length=max_length, | |
| temperature=temperature, | |
| top_p=top_p | |
| ) | |
| # --- Time Measurement End --- | |
| end_time = time.time() | |
| duration = round(end_time - start_time, 2) | |
| # For display purposes | |
| 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.") | |
| # Return all results including time and tokens | |
| return display_prompt, response, f"{duration} seconds", generated_tokens | |
| # --- System Prompts (Defaults only, UI will hold the editable version) --- | |
| 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" | |
| Base your analysis strictly on the provided frequencies and the stated RPS rules. | |
| """ | |
| # *** UPDATED Markov System Prompt v2 *** | |
| DEFAULT_SYSTEM_PROMPT_MARKOV = """You are an RPS assistant using short-term pattern analysis (Markov prediction). | |
| Your ONLY task is to recommend the AI move that beats the player's PREDICTED next move. Accuracy is critical. | |
| Input Information Provided: | |
| - Player's Predicted Next Move (from Markov analysis): [This is the key input!] | |
| Instructions: | |
| 1. **Identify Prediction:** State the player's PREDICTED next move (Rock, Paper, or Scissors) based *only* on the 'Predicted Next Move' input. | |
| 2. **Find Counter:** Apply the RPS rules (Paper beats Rock, Rock beats Scissors, Scissors beats Paper). Determine the single move that correctly beats the PREDICTED move from Step 1. State *only* the name of this counter move. Double-check the rules. | |
| 3. **Recommend:** Clearly state the counter move found in Step 2 as the AI's recommended move. | |
| Example Output Format: | |
| 1. Predicted Player Move: [Predicted move name] | |
| 2. Counter Move: [Counter move name] | |
| 3. Recommendation: Play [Counter move name]. | |
| """ | |
| # --- Default Input Values --- | |
| DEFAULT_PLAYER_STATS = "Rock: 40%\nPaper: 30%\nScissors: 30%" | |
| 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_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 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 | |
| ) | |
| # --- Visible System Prompt Textbox --- | |
| system_prompt_input = gr.Textbox( | |
| label="System Prompt (Edit based on selected mode)", | |
| value=DEFAULT_SYSTEM_PROMPT_FREQ, # Start with frequency prompt | |
| lines=15 | |
| ) | |
| # 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." | |
| ) | |
| with gr.Group(visible=False) as markov_inputs: # Hidden by default | |
| gr.Markdown("### Markov Prediction Analysis Inputs") | |
| player_last_move_input = gr.Dropdown( | |
| 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%)')." | |
| ) | |
| # General Inputs / Parameters / Outputs | |
| with gr.Row(): | |
| 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") | |
| 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) | |
| with gr.Column(): | |
| gr.Markdown(""" | |
| #### Testing Tips | |
| - Select the desired **Analysis Mode**. | |
| - Fill in the inputs for the **selected mode only**. | |
| - **Edit the System Prompt** above as needed for testing. | |
| - Use low **Temperature** for factual analysis. | |
| """) | |
| with gr.Row(): | |
| final_prompt_display = gr.Textbox( | |
| label="Formatted Input Sent to Model (via Chat Template)", lines=20 | |
| ) | |
| response_display = gr.Textbox( | |
| label="Model Response", lines=20, show_copy_button=True | |
| ) | |
| # --- Event Handlers --- | |
| # Function to update UI visibility AND system prompt content based on mode selection | |
| def update_ui_visibility_and_prompt(mode): | |
| if mode == "Frequency Only": | |
| return { | |
| frequency_inputs: gr.update(visible=True), | |
| markov_inputs: gr.update(visible=False), | |
| system_prompt_input: gr.update(value=DEFAULT_SYSTEM_PROMPT_FREQ) # Load Frequency prompt | |
| } | |
| elif mode == "Markov Prediction Only": | |
| return { | |
| frequency_inputs: gr.update(visible=False), | |
| markov_inputs: gr.update(visible=True), | |
| system_prompt_input: gr.update(value=DEFAULT_SYSTEM_PROMPT_MARKOV) # Load Markov prompt | |
| } | |
| else: # Default case | |
| return { | |
| frequency_inputs: gr.update(visible=True), | |
| markov_inputs: gr.update(visible=False), | |
| system_prompt_input: gr.update(value=DEFAULT_SYSTEM_PROMPT_FREQ) | |
| } | |
| # Link the radio button change to the UI update function | |
| analysis_mode_selector.change( | |
| fn=update_ui_visibility_and_prompt, # Use the combined update function | |
| inputs=analysis_mode_selector, | |
| outputs=[frequency_inputs, markov_inputs, system_prompt_input] # Components to update | |
| ) | |
| # Handle button click - Pass the single visible system prompt | |
| submit_btn.click( | |
| process_input, | |
| inputs=[ | |
| analysis_mode_selector, | |
| player_stats_input, | |
| player_last_move_input, | |
| markov_prediction_input, | |
| system_prompt_input, # Pass the visible system prompt textbox | |
| 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_selectable_analysis_v2" # Updated api_name | |
| ) | |
| # --- Launch the demo --- | |
| if __name__ == "__main__": | |
| demo.launch() | |