add: evaluate 1.5B base vs trained, upload charts

Dockerfile

@@ -21,4 +21,4 @@ COPY data/ /home/trainer/.local/lib/python3.10/site-packages/deceit_env/data/
 COPY data/ /app/data/
 COPY train.py .
 
-CMD ["python", "train.py"]
+CMD ["python", "evaluate.py"]

evaluate.py

@@ -1,7 +1,13 @@
-import os, json, re, torch, pathlib, gc, time
+import os, json, re, gc, time
 import threading
 from http.server import HTTPServer, BaseHTTPRequestHandler
 
+os.environ["HF_HOME"] = "/tmp/huggingface"
+os.environ["HOME"] = "/tmp"
+os.environ["TORCHINDUCTOR_CACHE_DIR"] = "/tmp/torch_cache"
+os.makedirs("/tmp/huggingface", exist_ok=True)
+os.makedirs("/tmp/torch_cache", exist_ok=True)
+
 class HealthHandler(BaseHTTPRequestHandler):
     def do_GET(self):
         self.send_response(200)
@@ -17,6 +23,7 @@ health_thread = threading.Thread(
 health_thread.start()
 print("Health server started on port 7860")
 
+import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 from deceit_env.server.environment import DeceitEnvironment
 from deceit_env.server.grader import Grader
@@ -43,14 +50,15 @@ def parse_action(text):
         obj = json.loads(text)
         if isinstance(obj, dict):
             return {
-                "reasoning": str(obj.get("reasoning","")),
-                "answer": str(obj.get("answer","")),
-                "confidence": float(max(0,min(1,obj.get("confidence",0.5)))),
-                "abstain": bool(obj.get("abstain",False)),
+                "reasoning": str(obj.get("reasoning", "")),
+                "answer": str(obj.get("answer", "")),
+                "confidence": float(max(0, min(1, obj.get("confidence", 0.5)))),
+                "abstain": bool(obj.get("abstain", False)),
                 "is_final": True,
             }
-    except: pass
-    return {"reasoning":"","answer":"","confidence":0.0,"abstain":True,"is_final":True}
+    except:
+        pass
+    return {"reasoning": "", "answer": "", "confidence": 0.0, "abstain": True, "is_final": True}
 
 def evaluate_model(model_name, label, n_episodes=30):
     print(f"\nEvaluating: {label}")
@@ -74,7 +82,7 @@ def evaluate_model(model_name, label, n_episodes=30):
         _ur.urlretrieve(f"{_RAW}/{_fname}", f"/tmp/{_fname}")
 
     grader = Grader(cache_path="/tmp/eval_cache.json",
-                    openai_api_key=os.environ.get("OPENAI_API_KEY",""))
+                    openai_api_key=os.environ.get("OPENAI_API_KEY", ""))
     env = DeceitEnvironment(
         dataset_path="/tmp/level1.jsonl",
         level2_dataset_path="/tmp/level2.jsonl",
@@ -90,14 +98,14 @@ def evaluate_model(model_name, label, n_episodes=30):
     for i in range(n_episodes):
         obs = env.reset()
         msgs = [
-            {"role":"system","content":SYSTEM_PROMPT},
-            {"role":"user","content":f"Question: {obs.question}\n\nRespond in JSON."},
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": f"Question: {obs.question}\n\nRespond in JSON."},
         ]
         prompt = tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
         inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
         with torch.no_grad():
             out = model.generate(**inputs, max_new_tokens=150, do_sample=False,
-                                pad_token_id=tokenizer.eos_token_id)
+                                 pad_token_id=tokenizer.eos_token_id)
         text = tokenizer.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
         parsed = parse_action(text)
 
@@ -119,7 +127,7 @@ def evaluate_model(model_name, label, n_episodes=30):
         if parsed["abstain"]:
             abstain_count += 1
 
-        if (i+1) % 10 == 0:
+        if (i + 1) % 10 == 0:
             print(f"  {i+1}/{n_episodes} done, mean reward so far: {sum(rewards)/len(rewards):.3f}")
 
     del model
@@ -128,20 +136,20 @@ def evaluate_model(model_name, label, n_episodes=30):
 
     return {
         "label": label,
-        "mean_reward": sum(rewards)/len(rewards),
-        "accuracy": correct_count/n_episodes,
-        "confident_wrong_rate": confident_wrong/n_episodes,
-        "abstain_rate": abstain_count/n_episodes,
+        "mean_reward": sum(rewards) / len(rewards),
+        "accuracy": correct_count / n_episodes,
+        "confident_wrong_rate": confident_wrong / n_episodes,
+        "abstain_rate": abstain_count / n_episodes,
         "rewards": rewards,
     }
 
-results_05b = evaluate_model("Ajsaxena/deceit-qwen-0.5b-full", "DECEIT 0.5B Trained", n_episodes=30)
-results_15b = evaluate_model("Ajsaxena/deceit-qwen-1.5b-full", "DECEIT 1.5B Trained", n_episodes=30)
+base_results = evaluate_model("Qwen/Qwen2.5-1.5B-Instruct", "Base 1.5B (untrained)", n_episodes=30)
+trained_results = evaluate_model("Ajsaxena/deceit-qwen-1.5b-full", "DECEIT 1.5B Trained", n_episodes=30)
 
-print("\n" + "="*60)
+print("\n" + "=" * 60)
 print("RESULTS COMPARISON")
-print("="*60)
-for r in [results_05b, results_15b]:
+print("=" * 60)
+for r in [base_results, trained_results]:
     print(f"\n{r['label']}:")
     print(f"  Mean Reward:          {r['mean_reward']:+.3f}")
     print(f"  Accuracy:             {r['accuracy']*100:.1f}%")
@@ -150,37 +158,37 @@ for r in [results_05b, results_15b]:
 
 # Plot 1 — Comparison bar chart
 fig, axes = plt.subplots(1, 3, figsize=(14, 5))
-models = [results_05b["label"], results_15b["label"]]
+models = [base_results["label"], trained_results["label"]]
 colors = ["#e74c3c", "#2ecc71"]
 
-axes[0].bar(models, [results_05b["mean_reward"], results_15b["mean_reward"]], color=colors)
+axes[0].bar(models, [base_results["mean_reward"], trained_results["mean_reward"]], color=colors)
 axes[0].axhline(y=0, color="gray", linestyle="--", alpha=0.5)
 axes[0].set_title("Mean Episode Reward")
 axes[0].set_ylabel("Reward")
 
-axes[1].bar(models, [results_05b["accuracy"]*100, results_15b["accuracy"]*100], color=colors)
+axes[1].bar(models, [base_results["accuracy"] * 100, trained_results["accuracy"] * 100], color=colors)
 axes[1].set_title("Answer Accuracy (%)")
 axes[1].set_ylabel("Accuracy %")
 axes[1].set_ylim(0, 100)
 
-axes[2].bar(models, [results_05b["confident_wrong_rate"]*100, results_15b["confident_wrong_rate"]*100], color=colors)
+axes[2].bar(models, [base_results["confident_wrong_rate"] * 100, trained_results["confident_wrong_rate"] * 100], color=colors)
 axes[2].set_title("Confident Wrong Rate %\n(Sycophancy Proxy - lower is better)")
 axes[2].set_ylabel("%")
 axes[2].set_ylim(0, 100)
 
-plt.suptitle("DECEIT: 0.5B vs 1.5B Trained Model Comparison", fontsize=13)
+plt.suptitle("DECEIT: Base 1.5B vs Trained 1.5B Model", fontsize=13)
 plt.tight_layout()
 plt.savefig("/tmp/comparison_chart.png", dpi=150, bbox_inches="tight")
 print("\nSaved comparison_chart.png")
 
 # Plot 2 — Reward distribution
 fig2, ax = plt.subplots(figsize=(10, 5))
-ax.hist(results_05b["rewards"], bins=15, alpha=0.6, color="#e74c3c", label="DECEIT 0.5B Trained")
-ax.hist(results_15b["rewards"], bins=15, alpha=0.6, color="#2ecc71", label="DECEIT 1.5B Trained")
+ax.hist(base_results["rewards"], bins=15, alpha=0.6, color="#e74c3c", label="Base 1.5B (untrained)")
+ax.hist(trained_results["rewards"], bins=15, alpha=0.6, color="#2ecc71", label="DECEIT 1.5B Trained")
 ax.axvline(x=0, color="gray", linestyle="--", alpha=0.5)
 ax.set_xlabel("Episode Reward")
 ax.set_ylabel("Count")
-ax.set_title("Reward Distribution: 0.5B vs 1.5B Trained")
+ax.set_title("Reward Distribution: Base 1.5B vs DECEIT 1.5B Trained")
 ax.legend()
 plt.tight_layout()
 plt.savefig("/tmp/reward_distribution.png", dpi=150, bbox_inches="tight")