Synced repo using 'sync_with_huggingface' Github Action

random_sample/arena_interface.py

@@ -6,14 +6,17 @@ from dotenv import load_dotenv
 load_dotenv()
 
 from .gen_api_answer import (
-    get_atla_response
+    get_atla_response,
+    get_selene_mini_response,
+    parse_selene_mini_response
 )
 
 from .prompts import (
     DEFAULT_EVAL_CRITERIA,
     DEFAULT_EVAL_PROMPT,
     DEFAULT_EVAL_PROMPT_EDITABLE,
-    FIXED_EVAL_SUFFIX
+    ATLA_PROMPT,
+    ATLA_PROMPT_WITH_REFERENCE
 )
 
 from .random_sample_generation import (
@@ -67,6 +70,15 @@ def create_arena_interface():
             value=DEFAULT_EVAL_PROMPT,
             visible=False
         )
+        with gr.Row():
+            # Add model selector dropdown at the top
+            model_selector = gr.Dropdown(
+                choices=["Selene", "Selene Mini"],
+                value="Selene",
+                label="Choose your Atla Model",
+                interactive=True
+            )
+
         with gr.Row():
             # Left side - Input section
             with gr.Column(scale=1):
@@ -234,36 +246,62 @@ def create_arena_interface():
         # Add a new state variable to track first game
         first_game_state = gr.State(True)  # Initialize as True
 
-        # Update the submit function to parse the evaluation criteria
+        # Update the submit function to handle both models
         def submit_and_store(
+            model_choice,
             use_reference,
             eval_criteria_text,
             human_input,
             ai_response,
-            ground_truth_input,
+            ground_truth,
         ):
-            # Build prompt data dictionary
-            prompt_data = {
-                'human_input': human_input,
-                'ai_response': ai_response,
-                'ground_truth_input': ground_truth_input if use_reference else None,
-                'eval_criteria': eval_criteria_text,
-            }
-
-            # Get response from Atla
-            response = get_atla_response(
-                model_name="AtlaAI/Selene-1-Mini-Llama-3.1-8B",
-                prompt=prompt_data,
-                max_tokens=500,
-                temperature=0.01
-            )
+            if model_choice == "Selene Mini":
+                # Prepare prompt based on reference mode
+                prompt_template = ATLA_PROMPT_WITH_REFERENCE if use_reference else ATLA_PROMPT
+                prompt = prompt_template.format(
+                    human_input=human_input,
+                    ai_response=ai_response,
+                    eval_criteria=eval_criteria_text,
+                    ground_truth=ground_truth if use_reference else ""
+                )
+                
+                print("\n=== Debug: Prompt being sent to Selene Mini ===")
+                print(prompt)
+                print("============================================\n")
+                
+                # Get and parse response
+                raw_response = get_selene_mini_response(
+                    model_name="AtlaAI/Selene-1-Mini-Llama-3.1-8B",
+                    prompt=prompt,
+                    max_tokens=500,
+                    temperature=0.01
+                )
+                response = parse_selene_mini_response(raw_response)
+            else:
+                # Selene API logic
+                prompt_data = {
+                    'human_input': human_input,
+                    'ai_response': ai_response,
+                    'ground_truth': ground_truth if use_reference else None,
+                    'eval_criteria': eval_criteria_text,
+                }
+                
+                print("\n=== Debug: Prompt data being sent to Selene API ===")
+                print(json.dumps(prompt_data, indent=2))
+                print("============================================\n")
+                
+                response = get_atla_response(
+                    model_name="AtlaAI/Selene-1-Mini-Llama-3.1-8B",
+                    prompt=prompt_data,
+                    max_tokens=500,
+                    temperature=0.01
+                )
 
             # Response now contains score and critique directly
             if isinstance(response, dict) and 'score' in response and 'critique' in response:
                 score = str(response['score'])
                 critique = response['critique']
             else:
-                # Handle error case
                 score = "Error"
                 critique = str(response)
 
@@ -274,22 +312,11 @@ def create_arena_interface():
                 gr.update(value="🎲"),
             ]
 
-        # Update the click handler to use False for is_first_game after first submission
-        def create_submit_handler():
-            first_game = True
-            
-            def handler(*args):
-                nonlocal first_game
-                result = submit_and_store(*args)
-                first_game = False  # Set to False after first submission
-                return result
-            
-            return handler
-
-        # Update the send_btn click handler
+        # Update the send_btn click handler with new input
         send_btn.click(
             fn=submit_and_store,
             inputs=[
+                model_selector,
                 use_reference_toggle,
                 eval_criteria_text,
                 human_input,

random_sample/gen_api_answer.py

@@ -1,14 +1,16 @@
 from openai import OpenAI
 import anthropic
 from together import Together
-import os
+import os       
 from atla import Atla
 from dotenv import load_dotenv
 from .prompts import (
-    JUDGE_SYSTEM_PROMPT,
-    ATLA_PROMPT,
-    ATLA_PROMPT_WITH_REFERENCE
+    JUDGE_SYSTEM_PROMPT
 )
+from transformers import AutoTokenizer
+import requests
+import json
+import re
 
 load_dotenv()
 
@@ -57,7 +59,7 @@ def get_atla_response(model_name, prompt, system_prompt=None, max_tokens=500, te
         # Extract components from the prompt data
         model_input = prompt.get('human_input', '')
         model_output = prompt.get('ai_response', '')
-        expected_output = prompt.get('ground_truth_input')
+        expected_output = prompt.get('ground_truth')
         evaluation_criteria = prompt.get('eval_criteria', '')
 
         response = atla_client.evaluation.create(
@@ -74,4 +76,73 @@ def get_atla_response(model_name, prompt, system_prompt=None, max_tokens=500, te
             "critique": response.result.evaluation.critique
         }
     except Exception as e:
-        return f"Error with Atla model {model_name}: {str(e)}"
+        return f"Error with Atla model {model_name}: {str(e)}"
+
+def get_selene_mini_response(model_name, prompt, system_prompt=None, max_tokens=500, temperature=0.01):
+    """Get response from HF endpoint for Atla model"""
+    try:
+        headers = {
+            "Accept": "application/json",
+            "Authorization": f"Bearer {hf_api_key}",
+            "Content-Type": "application/json"
+        }
+        
+        # Create messages list for chat template
+        messages = []
+        if system_prompt:
+            messages.append({"role": "system", "content": system_prompt})
+        messages.append({"role": "user", "content": prompt})
+        
+        # Apply chat template
+        model_id = "AtlaAI/Selene-1-Mini-Llama-3.1-8B"
+        tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_api_key)
+        formatted_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        
+        payload = {
+            "inputs": formatted_prompt,
+            "parameters": {
+                "max_new_tokens": max_tokens,
+                "return_full_text": False,
+                "temperature": temperature,
+                "seed": 42,
+                "add_generation_prompt": True
+            }
+        }
+        
+        response = requests.post(
+            "https://bkp9p28gri93egqh.us-east-1.aws.endpoints.huggingface.cloud",
+            headers=headers,
+            json=payload
+        )
+        return response.json()[0]["generated_text"]
+    except Exception as e:
+        return f"Error with Atla model {model_name}: {str(e)}"
+
+def parse_selene_mini_response(response_text):
+    """Parse the response from Selene Mini to extract score and critique"""
+    try:
+        # Clean up the response text
+        response_text = response_text.strip()
+        
+        # More flexible regex patterns
+        reasoning_pattern = r'\*\*Reasoning:?\*\*\s*(.*?)(?=\*\*Result|$)'
+        result_pattern = r'\*\*Result:?\*\*\s*(\d+)'
+        
+        reasoning_match = re.search(reasoning_pattern, response_text, re.DOTALL | re.IGNORECASE)
+        result_match = re.search(result_pattern, response_text, re.IGNORECASE)
+        
+        if reasoning_match and result_match:
+            critique = reasoning_match.group(1).strip()
+            score = result_match.group(1)
+            return {"score": score, "critique": critique}
+        else:
+            # If we can't parse it properly, let's return the raw response as critique
+            return {
+                "score": "Error",
+                "critique": f"Failed to parse response. Raw response:\n{response_text}"
+            }
+    except Exception as e:
+        return {
+            "score": "Error",
+            "critique": f"Error parsing response: {str(e)}\nRaw response:\n{response_text}"
+        }

random_sample/prompts.py

@@ -88,7 +88,7 @@ ATLA_PROMPT_WITH_REFERENCE = """You are tasked with evaluating a response based
   {eval_criteria}
 
   Reference answer:
-  {ground_truth_input}"""
+  {ground_truth}"""
 
 # Judge system prompt for non-Prometheus models
 JUDGE_SYSTEM_PROMPT = """Please act as an impartial judge and evaluate based on the user's instruction. Your output format should strictly adhere to JSON as follows: {"feedback": "<write feedback>", "result": <numerical score>}. Ensure the output is valid JSON, without additional formatting or explanations."""