Update app.py

Added Judge model

app.py

@@ -128,12 +128,49 @@ def benchmark(config_text, dataset_text, task):
 
     return df, "\n".join(summary)
 
+
+def hf_judge(model_name, prompt, candidate, reference=None, rubric=None, max_new_tokens=256):
+    """
+    Calls a judge model to score candidate output. Returns a dict of scores.
+    Rubric should instruct JSON output, e.g.:
+    {"relevance": int, "factuality": int, "clarity": int, "overall": float}
+    """
+    from huggingface_hub import InferenceClient
+    client = InferenceClient(model=model_name, token=os.getenv("HF_TOKEN"))
+
+    rubric = rubric or (
+        "Evaluate the candidate answer. Score 1–5 for:\n"
+        "- Relevance: addresses the prompt\n"
+        "- Factuality: correct and supported\n"
+        "- Clarity: clear and well-structured\n"
+        "Return JSON: {\"relevance\": int, \"factuality\": int, \"clarity\": int, \"overall\": float}"
+    )
+
+    judge_prompt = (
+        f"{rubric}\n\nPrompt:\n{prompt}\n\n"
+        f"Candidate:\n{candidate}\n\n"
+        f"Reference (if available):\n{reference if reference is not None else 'N/A'}\n"
+    )
+
+    try:
+        text = client.text_generation(judge_prompt, max_new_tokens=max_new_tokens, temperature=0.0)
+        # Attempt to parse JSON anywhere in the response:
+        import re, json
+        m = re.search(r'\{.*\}', text, re.S)
+        return json.loads(m.group(0)) if m else {"raw": text}
+    except Exception as e:
+        return {"error": str(e)}
+
+
 # ---------------- Gradio UI ---------------- #
 with gr.Blocks() as demo:
     gr.Markdown("# LLM Benchmarking App (Hugging Face)")
     gr.Markdown("Upload config.yaml and dataset.jsonl, select task, and run benchmark.")
+    
 
     with gr.Row():
+        use_judge = gr.Checkbox(label="Use judge model", value=False)
+        judge_model = gr.Textbox(label="Judge model (HF repo id)", value="mistralai/Mistral-7B-Instruct")
         config_file = gr.File(label="Upload Config (YAML)", type="filepath")
         dataset_file = gr.File(label="Upload Dataset (JSONL)", type="filepath")
     task = gr.Textbox(label="Task (e.g., qa, summarization, classification)")
@@ -142,15 +179,59 @@ with gr.Blocks() as demo:
     results_table = gr.Dataframe(headers=["model","id","task","prediction","reference","latency_seconds","exact_match","f1","rougeL_f","bleu"], label="Results")
     summary_box = gr.Textbox(label="Summary", lines=10)
     download_csv = gr.File(label="Download CSV")
+    
+    
 
-    def run_benchmark(config_path, dataset_path, task):
+
+    
+    def run_benchmark(config_path, dataset_path, task, use_judge=False, judge_model=None):
         config_text = open(config_path, "r", encoding="utf-8").read()
         dataset_text = open(dataset_path, "r", encoding="utf-8").read()
-        df, summary = benchmark(config_text, dataset_text, task)
-        csv_path = "results.csv"
-        df.to_csv(csv_path, index=False)
-        return df, summary, csv_path
-
-    run_btn.click(run_benchmark, inputs=[config_file, dataset_file, task], outputs=[results_table, summary_box, download_csv])
+        cfg = yaml.safe_load(config_text)
+        data = [json.loads(line) for line in dataset_text.splitlines() if line.strip()]
+        template = cfg.get("prompt_templates", {}).get(task, "{{text}}")
+    
+        results = []
+        for m in cfg.get("models", []):
+            model_name = m["name"]
+            max_tokens = m.get("params", {}).get("max_tokens", 256)
+            temperature = m.get("params", {}).get("temperature", 0.2)
+    
+            for ex in data:
+                variables = {k: ex.get(k, "") for k in ("question", "context", "text", "labels")}
+                prompt = template
+                for k, v in variables.items():
+                    prompt = prompt.replace(f"{{{{{k}}}}}", str(v))
+    
+                prediction, latency = hf_generate(model_name, prompt, max_new_tokens=max_tokens, temperature=temperature)
+                metrics = compute_metrics(task, prediction, ex.get("reference", ""))
+    
+                row = {
+                    "model": model_name,
+                    "id": ex.get("id", ""),
+                    "task": task,
+                    "prompt": prompt,
+                    "prediction": prediction,
+                    "reference": ex.get("reference", ""),
+                    "latency_seconds": latency,
+                    **metrics
+                }
+    
+                if use_judge and judge_model:
+                    scores = hf_judge(judge_model, prompt, prediction, ex.get("reference", ""))
+                    # Flatten judge scores with prefix
+                    for k, v in (scores.items() if isinstance(scores, dict) else []):
+                        row[f"judge_{k}"] = v
+    
+                results.append(row)
+    
+        # Save results and summary as before …
+        # (Then return df, summary, csv_path)
+        
+        run_btn.click(
+            run_benchmark,
+            inputs=[config_file, dataset_file, task, use_judge, judge_model],
+            outputs=[results_table, summary_box, download_csv]
+        )
 
 demo.launch()