feat: add flow judge model

app.py

@@ -13,7 +13,8 @@ import gradio as gr
 from gen_api_answer import (
     get_model_response, 
     parse_model_response,
-    prometheus_parse_model_response
+    prometheus_parse_model_response,
+    flow_judge_parse_model_response,
 )
 
 from random_sample_generation import (
@@ -749,10 +750,17 @@ with gr.Blocks(theme="default", css=CSS_STYLES) as demo:
         # Parse the responses based on model, using Prometheus parsing for Prometheus models and JSON parsing for others
         is_prometheus_a = (model_data.get(model_a)['organization'] == 'Prometheus')
         is_prometheus_b = (model_data.get(model_b)['organization'] == 'Prometheus')
+        
+        # Parse the responses based on model, using Flow Judge parsing for Flow Judge models and Prometheus parsing for others
+        is_flow_judge_a = (model_data.get(model_a)['organization'] == 'Flow AI')
+        is_flow_judge_b = (model_data.get(model_b)['organization'] == 'Flow AI')
 
         if is_prometheus_a:
             score_a_val, critique_a_val = prometheus_parse_model_response(response_a)
             score_a_val = f"{score_a_val} / 5"
+        elif is_flow_judge_a:
+            score_a_val, critique_a_val = flow_judge_parse_model_response(response_a)
+            score_a_val = f"{score_a_val} / 5"
         else:
             score_a_val, critique_a_val = parse_model_response(response_a)
             score_a_val = f"{score_a_val} / 5"
@@ -760,6 +768,9 @@ with gr.Blocks(theme="default", css=CSS_STYLES) as demo:
         if is_prometheus_b:
             score_b_val, critique_b_val = prometheus_parse_model_response(response_b)
             score_b_val = f"{score_b_val} / 5"
+        elif is_flow_judge_b:
+            score_b_val, critique_b_val = flow_judge_parse_model_response(response_b)
+            score_b_val = f"{score_b_val} / 5"
         else:
             score_b_val, critique_b_val = parse_model_response(response_b)
             score_b_val = f"{score_b_val} / 5"

data/models.jsonl

@@ -18,4 +18,5 @@
 {"name": "Claude 3.5 Haiku", "organization": "Anthropic", "license": "Proprietary", "api_model": "claude-3-5-haiku-latest"}
 {"name": "Prometheus-7b v2", "organization": "Prometheus", "license": "Open Source", "api_model": "prometheus/prometheus-7b-v2"}
 {"name": "Command-R", "organization": "Cohere", "license": "Proprietary", "api_model": "command-r"}
-{"name": "Command-R Plus", "organization": "Cohere", "license": "Proprietary", "api_model": "command-r-plus"}
+{"name": "Command-R Plus", "organization": "Cohere", "license": "Proprietary", "api_model": "command-r-plus"}
+{"name": "Flow-Judge-v0.1", "organization": "Flow AI", "license": "Open Source", "api_model": "Flow-Judge-v0.1-4.65bpw-exl2"}

gen_api_answer.py

@@ -10,6 +10,7 @@ from prompts import (
     JUDGE_SYSTEM_PROMPT,
     PROMETHEUS_PROMPT,
     PROMETHEUS_PROMPT_WITH_REFERENCE,
+    FLOW_JUDGE_PROMPT
 )
 
 # Initialize clients
@@ -22,6 +23,8 @@ huggingface_client = OpenAI(
     base_url="https://otb7jglxy6r37af6.us-east-1.aws.endpoints.huggingface.cloud/v1/",
     api_key=hf_api_key
 )
+flow_judge_api_key = os.getenv("FLOW_JUDGE_API_KEY")
+
 
 def get_openai_response(model_name, prompt, system_prompt=JUDGE_SYSTEM_PROMPT, max_tokens=500, temperature=0):
     """Get response from OpenAI API"""
@@ -116,6 +119,30 @@ def get_cohere_response(model_name, prompt, system_prompt=JUDGE_SYSTEM_PROMPT, m
         return str(content_items)  # Fallback if it's not a list
     except Exception as e:
         return f"Error with Cohere model {model_name}: {str(e)}"
+    
+def get_flow_judge_response(model_name, prompt, max_tokens=500, temperature=0.1, top_p=0.95) -> str:
+    """Get response from Flow Judge"""
+    try:
+        response = requests.post(
+            "https://tsukuyomi.tailfa581.ts.net/v1/chat/completions",
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {flow_judge_api_key}"
+            },
+            json={
+                "model": model_name,
+                "messages": [
+                    {"role": "user", "content": prompt}
+                ],
+                "max_tokens": max_tokens,
+                "temperature": temperature,
+                "top_p": top_p
+            }
+        )
+        response.raise_for_status()
+        return response.json()["choices"][0]['message']['content']
+    except Exception as e:
+        return f"Error with Flow Judge completions model {model_name}: {str(e)}"
 
 def get_model_response(
     model_name,
@@ -134,36 +161,68 @@ def get_model_response(
 
     # Determine if model is Prometheus
     is_prometheus = (organization == "Prometheus")
-
+    is_flow_judge = (organization == "Flow AI")
     # For non-Prometheus models, use the Judge system prompt
-    system_prompt = None if is_prometheus else JUDGE_SYSTEM_PROMPT
+    system_prompt = None if is_prometheus or is_flow_judge else JUDGE_SYSTEM_PROMPT
 
     # Select the appropriate base prompt
     if use_reference:
-        base_prompt = PROMETHEUS_PROMPT_WITH_REFERENCE
+        if not is_flow_judge:
+            base_prompt = PROMETHEUS_PROMPT_WITH_REFERENCE
+        else:
+            base_prompt = FLOW_JUDGE_PROMPT
     else:
-        base_prompt = PROMETHEUS_PROMPT
+        if not is_flow_judge:
+            base_prompt = PROMETHEUS_PROMPT
+        else:
+            base_prompt = FLOW_JUDGE_PROMPT
 
     # For non-Prometheus models, replace the specific instruction
-    if not is_prometheus:
+    if not is_prometheus and not is_flow_judge:
         base_prompt = base_prompt.replace(
             '3. The output format should look as follows: "Feedback: (write a feedback for criteria) [RESULT] (an integer number between 1 and 5)"',
             '3. 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.'
         )
 
     try:
-        # Format the prompt with the provided data, only using available keys
-        final_prompt = base_prompt.format(
-            human_input=prompt_data['human_input'],
-            ai_response=prompt_data['ai_response'],
-            ground_truth_input=prompt_data.get('ground_truth_input', ''),
-            eval_criteria=prompt_data['eval_criteria'],
-            score1_desc=prompt_data['score1_desc'],
-            score2_desc=prompt_data['score2_desc'],
-            score3_desc=prompt_data['score3_desc'],
-            score4_desc=prompt_data['score4_desc'],
-            score5_desc=prompt_data['score5_desc']
-        )
+        if not is_flow_judge:
+            # Format the prompt with the provided data, only using available keys
+            final_prompt = base_prompt.format(
+                human_input=prompt_data['human_input'],
+                ai_response=prompt_data['ai_response'],
+                ground_truth_input=prompt_data.get('ground_truth_input', ''),
+                eval_criteria=prompt_data['eval_criteria'],
+                score1_desc=prompt_data['score1_desc'],
+                score2_desc=prompt_data['score2_desc'],
+                score3_desc=prompt_data['score3_desc'],
+                score4_desc=prompt_data['score4_desc'],
+                score5_desc=prompt_data['score5_desc']
+            )
+        else:
+            human_input = f"<user_input>\n{prompt_data['human_input']}\n</user_input>"
+            ai_response = f"<response>\n{prompt_data['ai_response']}\n</response>"
+            ground_truth=prompt_data.get('ground_truth_input', '')
+            if ground_truth:
+                response_reference = f"<response_reference>\n{ground_truth}\n</response_reference>"
+            else:
+                response_reference = ""
+            eval_criteria = prompt_data['eval_criteria']
+            score1_desc = f"- Score 1: {prompt_data['score1_desc']}\n"
+            score2_desc = f"- Score 2: {prompt_data['score2_desc']}\n"
+            score3_desc = f"- Score 3: {prompt_data['score3_desc']}\n"
+            score4_desc = f"- Score 4: {prompt_data['score4_desc']}\n"
+            score5_desc = f"- Score 5: {prompt_data['score5_desc']}"
+            rubric = score1_desc + score2_desc + score3_desc + score4_desc + score5_desc
+            if response_reference:
+                inputs = human_input + "\n"+ response_reference
+            else:
+                inputs = human_input
+            final_prompt = base_prompt.format(
+                INPUTS=inputs,
+                OUTPUT=ai_response,
+                EVALUATION_CRITERIA=eval_criteria,
+                RUBRIC=rubric
+            )
     except KeyError as e:
         return f"Error formatting prompt: Missing required field {str(e)}"
 
@@ -184,6 +243,10 @@ def get_model_response(
             return get_cohere_response(
                 api_model, final_prompt, system_prompt, max_tokens, temperature
             )
+        elif organization == "Flow AI":
+            return get_flow_judge_response(
+                api_model, final_prompt, max_tokens, temperature
+            )
         else:
             # All other organizations use Together API
             return get_together_response(
@@ -267,6 +330,30 @@ def prometheus_parse_model_response(output):
 
         return "Error", f"Failed to parse response: {output}"
 
+    except Exception as e:
+        print(f"Failed to parse response: {str(e)}")
+        return "Error", f"Exception during parsing: {str(e)}"
+    
+def flow_judge_parse_model_response(output):
+    try:
+        print(f"Raw model response: {output}")
+        # Convert multiple line breaks to single ones and strip whitespace
+        output = re.sub(r'\n{2,}', '\n', output.strip())
+        
+        # Compile regex patterns
+        feedback_pattern = re.compile(r"<feedback>\s*(.*?)\s*</feedback>", re.DOTALL)
+        score_pattern = re.compile(r"<score>\s*(\d+)\s*</score>", re.DOTALL)
+
+        feedback_match = feedback_pattern.search(output)
+        score_match = score_pattern.search(output)
+
+        if feedback_match or not score_match:
+            feedback = feedback_match.group(1).strip()
+            score = int(score_match.group(1).strip())
+            return str(score), feedback
+            
+        return "Error", f"Failed to parse response: {output}"
+        
     except Exception as e:
         print(f"Failed to parse response: {str(e)}")
         return "Error", f"Exception during parsing: {str(e)}"

prompts.py

@@ -90,5 +90,59 @@ Score 5: {score5_desc}
 ###Feedback:
 """
 
+# Define the Flow Judge prompt
+FLOW_JUDGE_PROMPT = """# GOAL
+Your job is to evaluate a task carried out by an AI system powered by a large \
+language model.
+
+You will be provided with the inputs and output of the task, as well as the evaluation criteria \
+and scoring rubric. Your task is to evaluate the output of the AI system based on the evaluation \
+criteria and scoring rubric provided.
+
+# INPUT
+Below are the inputs required for performing the task:
+<inputs>
+{INPUTS}
+</inputs>
+
+# OUTPUT
+Below is the output of the task:
+<output>
+{OUTPUT}
+</output>
+
+# EVALUATION CRITERIA AND SCORING RUBRIC
+Here are the evaluation criteria and the rubric that you need to use for evaluating the task:
+<evaluation_criteria>
+{EVALUATION_CRITERIA}
+</evaluation_criteria>
+
+<scoring_rubric>
+{RUBRIC}
+</scoring_rubric>
+
+# INSTRUCTIONS FOR THE EVALUATION
+1. Understand the task and criteria: Familiarize yourself with the task to be evaluated. \
+Review the evaluation criteria and scoring rubric to understand the different levels of \
+performance and the descriptions for each score.
+2. Review the inputs and output: Look at the inputs provided for the task. Examine the output \
+generated from completing the task.
+3. Compare output to score descriptions: Compare the output against the criteria and score \
+descriptions in the scoring rubric. For each criterion,decide which description best matches the \
+output.
+4. After comparing the output to the score descriptions, pay attention to the small details that \
+might impact the final score that you assign. Sometimes a small difference can dictate the final \
+score.
+5. Write verbal feedback justifying your evaluation that includes a detailed rationale, referring \
+to specific aspects of the output and comparing them to the rubric.
+6. Assign a final score based on the scoring rubric.
+
+## FORMAT FOR THE EVALUATION
+- Write the verbal feedback inside <feedback> tags without any additional surrounding text.
+- Write the numeric score inside <score> tags, without any additional surrounding text and always \
+after the feedback.
+
+Please accurately evaluate the task. Strictly adhere to the evaluation criteria and rubric."""
+
 # 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."""