Update app.py

app.py

@@ -2,6 +2,8 @@ import os
 import json
 import re
 import uuid
+import io
+import csv
 from datetime import datetime
 import openai
 import gradio as gr
@@ -31,12 +33,9 @@ GRADER_MODEL = "gpt-4o-mini"
 openai.api_key = os.getenv("OPENAI_API_KEY")
 genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
 
-# Models that only support default temperature (don’t allow setting temperature manually)
+# Models that only support default temperature
 MODEL_DEFAULT_TEMP = ["o4-mini"]
 
-# Local JSON file for storing runs
-RUNS_FILE = "/data/runs2.json"
-
 # -------------------------
 # Helper to read JSONL
 # -------------------------
@@ -67,9 +66,9 @@ class HealthBenchEval:
 
         self.scores = []
         self.htmls = ""
-        self.sample_records = []  # <-- store per-sample rows
+        self.sample_records = []
         self.seed = seed
-        self.eval_id = str(uuid.uuid4())  # unique id per evaluation run
+        self.eval_id = str(uuid.uuid4())
 
     def score_with_grader(self, prompt_text, completion_text, example_index):
         prompt = f"""
@@ -96,10 +95,7 @@ Return only a number between 0 and 1.
             return 0.0
 
     def generate_with_candidate(self, candidate_model, system_prompt, prompt_text, example_index, max_tokens=1024):
-        """
-        Generate completion with retry logic and better error logging.
-        """
-        for attempt in range(3):  # retry up to 3 times
+        for attempt in range(3):
             try:
                 if candidate_model.startswith("gemini"):
                     model = genai.GenerativeModel(candidate_model)
@@ -119,7 +115,6 @@ Return only a number between 0 and 1.
                         messages.append({"role": "system", "content": system_prompt})
                     messages.append({"role": "user", "content": prompt_text})
 
-                    # Skip temperature for models that don't support it
                     if candidate_model in MODEL_DEFAULT_TEMP:
                         resp = openai.chat.completions.create(
                             model=candidate_model,
@@ -134,7 +129,6 @@ Return only a number between 0 and 1.
                             max_completion_tokens=max_tokens
                         )
                     completion = resp.choices[0].message.content
-                    print(resp)
 
                 return completion.strip() if hasattr(completion, "strip") else completion
 
@@ -142,7 +136,7 @@ Return only a number between 0 and 1.
                 print(f"[ERROR] Candidate model {candidate_model} failed at dataset index {example_index} (attempt {attempt+1}/3)")
                 print(f"Prompt text: {prompt_text[:200]}...")
                 print(f"Error: {e}")
-                if attempt == 2:  # after last attempt
+                if attempt == 2:
                     return f"[ERROR after 3 retries: {str(e)}]"
 
     def __call__(self, candidate_model, system_prompt, eval_subset=""):
@@ -150,21 +144,19 @@ Return only a number between 0 and 1.
 
         cumulative_total = 0.0
         for i, example in enumerate(self.dataset):
-            dataset_index = self.indices[i]  # actual dataset row index
+            dataset_index = self.indices[i]
             prompt_obj = example.get("prompt", [])
             prompt_text = " ".join([m.get("content", "") for m in prompt_obj])
 
             completion_text = self.generate_with_candidate(candidate_model, system_prompt, prompt_text, dataset_index)
             score = self.score_with_grader(prompt_text, completion_text, dataset_index)
 
-            # update running totals (per eval_id)
             cumulative_total += score
             cumulative_avg = cumulative_total / (i + 1)
 
             self.scores.append(score)
             html_lines.append(f"<li>Dataset Row {dataset_index}: Score = {score:.3f}</li>")
 
-            # create individual sample record
             self.sample_records.append({
                 "eval_id": self.eval_id,
                 "timestamp": datetime.utcnow().isoformat(),
@@ -184,19 +176,9 @@ Return only a number between 0 and 1.
         return self
 
 # -------------------------
-# Helper to generate HTML table from runs
+# Helpers
 # -------------------------
-def generate_runs_html():
-    runs = []
-    if os.path.exists(RUNS_FILE):
-        try:
-            with open(RUNS_FILE, "r", encoding="utf-8") as f:
-                runs = json.load(f)
-                if not isinstance(runs, list):
-                    runs = []
-        except (json.JSONDecodeError, ValueError):
-            runs = []
-
+def generate_runs_html(runs):
     if runs:
         table_rows = ""
         for r in reversed(runs):
@@ -244,24 +226,40 @@ def generate_runs_html():
         """
     else:
         runs_html = "<p>No evaluations yet.</p>"
-
     return runs_html
 
-# -------------------------
-# Clear runs file
-# -------------------------
-def clear_runs():
-    with open(RUNS_FILE, "w", encoding="utf-8") as f:
-        json.dump([], f, indent=2)
-    return "<p>No evaluations yet.</p>"
+def generate_csv(runs):
+    if not runs:
+        return None
+    output = io.StringIO()
+    fieldnames = ["eval_id", "timestamp", "candidate_model", "system_prompt", "eval_subset", 
+                  "seed", "dataset_index", "prompt_text", "completion_text", "score", 
+                  "cumulative_total", "cumulative_avg"]
+    writer = csv.DictWriter(output, fieldnames=fieldnames)
+    writer.writeheader()
+    for run in runs:
+        writer.writerow(run)
+    csv_data = output.getvalue()
+    output.close()
+    return csv_data
+
+def prepare_download(runs):
+    csv_data = generate_csv(runs)
+    if not csv_data:
+        return None
+    filename = f"eval_results_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"
+    filepath = os.path.join("/tmp", filename)
+    with open(filepath, "w", encoding="utf-8") as f:
+        f.write(csv_data)
+    return filepath
 
 # -------------------------
 # Gradio UI function
 # -------------------------
-def run_eval_ui(candidate_model, system_prompt, eval_subset, num_examples, seed):
+def run_eval_ui(candidate_model, system_prompt, eval_subset, num_examples, seed, runs):
     dataset_file = DATASET_FILES.get(eval_subset)
     if not dataset_file:
-        return "<p style='color:red'>Invalid dataset</p>", {}, generate_runs_html()
+        return "<p style='color:red'>Invalid dataset</p>", {}, generate_runs_html(runs), runs
 
     seed_val = int(seed) if seed else None
     num_val = int(num_examples) if num_examples else None
@@ -269,23 +267,8 @@ def run_eval_ui(candidate_model, system_prompt, eval_subset, num_examples, seed)
     eval_obj = HealthBenchEval(dataset_file, num_examples=num_val, seed=seed_val)
     result = eval_obj(candidate_model, system_prompt, eval_subset=eval_subset)
 
-    # Load existing runs
-    runs = []
-    if os.path.exists(RUNS_FILE):
-        try:
-            with open(RUNS_FILE, "r", encoding="utf-8") as f:
-                runs = json.load(f)
-                if not isinstance(runs, list):
-                    runs = []
-        except (json.JSONDecodeError, ValueError):
-            runs = []
-
     runs.extend(result.sample_records)
-
-    with open(RUNS_FILE, "w", encoding="utf-8") as f:
-        json.dump(runs, f, indent=2)
-
-    runs_html = generate_runs_html()
+    runs_html = generate_runs_html(runs)
 
     metrics = {
         "eval_id": result.eval_id,
@@ -295,7 +278,10 @@ def run_eval_ui(candidate_model, system_prompt, eval_subset, num_examples, seed)
         "seed": seed_val
     }
 
-    return result.htmls, metrics, runs_html
+    return result.htmls, metrics, runs_html, runs
+
+def clear_runs():
+    return "<p>No evaluations yet.</p>", []
 
 # -------------------------
 # Gradio UI
@@ -328,22 +314,32 @@ def ui():
 
         output_html = gr.HTML(label="Evaluation Report")
         output_metrics = gr.JSON(label="Metrics JSON")
-        output_all_runs = gr.HTML(label="Evaluation History", value=generate_runs_html())
+        output_all_runs = gr.HTML(label="Evaluation History")
+        session_runs = gr.State([])
 
         with gr.Row():
             clear_btn = gr.Button("Clear History")
+            download_btn = gr.DownloadButton(
+                label="Download CSV",
+                variant="secondary"
+            )
 
-        # Connect buttons
         run_btn.click(
             fn=run_eval_ui,
-            inputs=[candidate_model, system_prompt, eval_subset, num_examples, seed],
-            outputs=[output_html, output_metrics, output_all_runs]
+            inputs=[candidate_model, system_prompt, eval_subset, num_examples, seed, session_runs],
+            outputs=[output_html, output_metrics, output_all_runs, session_runs]
         )
 
         clear_btn.click(
             fn=clear_runs,
             inputs=[],
-            outputs=[output_all_runs]
+            outputs=[output_all_runs, session_runs]
+        )
+
+        download_btn.click(
+            fn=prepare_download,
+            inputs=[session_runs],
+            outputs=[download_btn]
         )
 
     return demo