Update app.py

app.py

@@ -162,50 +162,57 @@ def hf_judge(model_name, prompt, candidate, reference=None, rubric=None, max_new
         return {"error": str(e)}
 
 
+# ---------------- Gradio UI ---------------- #
+
 # ---------------- 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)")
+    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")
+    rubric = gr.Textbox(label="Judge Rubric", value="Evaluate relevance, factuality, clarity. Return JSON.")
+
     run_btn = gr.Button("Run Benchmark")
 
-    results_table = gr.Dataframe(headers=["model","id","task","prediction","reference","latency_seconds","exact_match","f1","rougeL_f","bleu"], label="Results")
+    results_table = gr.Dataframe(headers=[
+        "model","id","task","prediction","reference","latency_seconds",
+        "exact_match","f1","rougeL_f","bleu","judge_relevance","judge_factuality","judge_clarity","judge_overall"
+    ], 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, use_judge=False, judge_model=None, rubric=None):
+        if not config_path or not dataset_path:
+            return None, "Error: Please upload both files", None
 
-    
-    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()
         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", ""),
@@ -216,22 +223,38 @@ with gr.Blocks() as demo:
                     "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
+                    scores = hf_judge(judge_model, prompt, prediction, ex.get("reference", ""), rubric)
                     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]
-        )
+
+        df = pd.DataFrame(results)
+        csv_path = "results.csv"
+        df.to_csv(csv_path, index=False)
+
+        # Summary
+        summary = []
+        for model_name in set(df["model"]):
+            sub = df[df["model"] == model_name]
+            summary.append(f"## {model_name}")
+            summary.append(f"Samples: {len(sub)}")
+            for metric in ["exact_match", "f1", "rougeL_f", "bleu", "judge_overall"]:
+                if metric in sub.columns:
+                    vals = [v for v in sub[metric] if isinstance(v, (int, float))]
+                    if vals:
+                        summary.append(f"{metric}: mean={sum(vals)/len(vals):.4f}")
+            summary.append(f"Latency mean: {sum(sub['latency_seconds'])/len(sub):.3f}s\n")
+
+        return df, "\n".join(summary), csv_path
+
+    run_btn.click(
+        run_benchmark,
+        inputs=[config_file, dataset_file, task, use_judge, judge_model, rubric],
+        outputs=[results_table, summary_box, download_csv]
+    )
 
 demo.launch()
+