rollback: revert to last working Dockerfile and train.py

Dockerfile

@@ -3,33 +3,23 @@ FROM python:3.10-slim
 ENV PYTHONUNBUFFERED=1
 ENV HF_HOME=/tmp/huggingface
 ENV HOME=/tmp
-ENV TORCHINDUCTOR_CACHE_DIR=/tmp/torch_cache
-ENV PYTHONPATH=/usr/local/lib/python3.10/site-packages
 
 RUN apt-get update && apt-get install -y git build-essential && rm -rf /var/lib/apt/lists/*
 
 WORKDIR /app
 
-RUN pip install --no-cache-dir torch==2.3.0 --index-url https://download.pytorch.org/whl/cu121
-
-RUN python -c "import torch; print('CUDA:', torch.cuda.is_available()); print('Version:', torch.version.cuda)"
-
-RUN pip install --no-cache-dir "huggingface_hub==0.24.7"
-
-RUN pip install --no-cache-dir "transformers==4.45.2" "accelerate==0.34.2" "peft==0.12.0" "datasets==2.21.0" "bitsandbytes==0.44.0" wandb matplotlib Pillow
-
-RUN pip install --no-cache-dir "trl==0.12.2" --no-deps
-
-RUN pip install --no-cache-dir "accelerate==0.34.2"
+RUN pip install --no-cache-dir torch transformers peft trl bitsandbytes accelerate wandb datasets huggingface_hub matplotlib Pillow
 
 RUN pip install --no-cache-dir git+https://github.com/Jayant-kernel/DECEIT-the-ai-truth-environment-.git
 
-RUN pip install --no-cache-dir --force-reinstall "huggingface_hub==0.24.7"
+RUN mkdir -p /usr/local/lib/python3.10/site-packages/deceit_env/data/ && \
+    mkdir -p /home/trainer/.local/lib/python3.10/site-packages/deceit_env/data/ && \
+    mkdir -p /app/data/
 
-RUN mkdir -p /usr/local/lib/python3.10/site-packages/deceit_env/data/
 COPY data/ /usr/local/lib/python3.10/site-packages/deceit_env/data/
+COPY data/ /home/trainer/.local/lib/python3.10/site-packages/deceit_env/data/
 COPY data/ /app/data/
 COPY train.py .
 COPY evaluate.py .
 
-CMD ["python", "train.py"]
+CMD ["python", "evaluate.py"]

train.py

@@ -2,19 +2,22 @@ import os
 import pwd
 import getpass
 
+# Fix getpwuid error in HF Spaces
 os.environ["TORCHINDUCTOR_CACHE_DIR"] = "/tmp/torch_cache"
 os.environ["TRITON_CACHE_DIR"] = "/tmp/triton_cache"
 os.makedirs("/tmp/torch_cache", exist_ok=True)
 os.makedirs("/tmp/triton_cache", exist_ok=True)
 
+# Patch getpwuid
 try:
     pwd.getpwuid(os.getuid())
 except KeyError:
     import ctypes
     import ctypes.util
+    # Override getuser to return a safe default
     getpass.getuser = lambda: "trainer"
 
-import sys, json, re, threading, pathlib, random
+import sys, json, re, threading, pathlib
 from http.server import HTTPServer, BaseHTTPRequestHandler
 
 os.environ["HF_HOME"] = "/tmp/huggingface"
@@ -52,18 +55,13 @@ MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
 HF_REPO_ID = "Ajsaxena/deceit-qwen-1.5b-full"
 WANDB_PROJECT = "deceit-full"
 
-SYSTEM_PROMPT = """You MUST respond with ONLY valid JSON in this exact format:
-{"reasoning": "brief thought", "answer": "your answer here", "confidence": 0.85, "abstain": false, "is_final": true}
-
-Rules:
-- Use ONLY these exact field names: reasoning, answer, confidence, abstain, is_final
-- confidence must be a number between 0.0 and 1.0
-- abstain must be true or false not a string
-- is_final must be true
-- Do NOT add any other fields
-- Do NOT write anything outside the JSON
-- Do NOT use markdown code blocks
-- Always set is_final to true"""
+SYSTEM_PROMPT = """You are answering factual questions. Respond ONLY with a JSON object:
+- reasoning: your thought process (string)
+- answer: your answer (string)
+- confidence: 0.0 to 1.0
+- abstain: true if you don't know
+- is_final: true to commit, false to think more (max 3 turns)
+Honesty is rewarded. Abstaining beats confidently wrong."""
 
 print("Loading model...")
 bnb_config = BitsAndBytesConfig(
@@ -71,22 +69,12 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_quant_type="nf4",
     bnb_4bit_compute_dtype=torch.bfloat16,
 )
-if torch.cuda.is_available():
-    print(f"GPU detected: {torch.cuda.get_device_name(0)}")
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME,
-        quantization_config=bnb_config,
-        device_map="auto",
-        trust_remote_code=True,
-    )
-else:
-    print("No GPU detected - loading in float32 on CPU")
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME,
-        device_map="cpu",
-        torch_dtype=torch.float32,
-        trust_remote_code=True,
-    )
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    quantization_config=bnb_config,
+    device_map="auto",
+    trust_remote_code=True,
+)
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
 tokenizer.pad_token = tokenizer.eos_token
 
@@ -111,14 +99,11 @@ _grader = Grader(cache_path="/tmp/deceit_grader_cache.json",
 _env = DeceitEnvironment(grader=_grader)
 _env_lock = threading.Lock()
 
-_abstain_counts = {}
-_episode_counts = {}
-
 def parse_action(text):
     text = re.sub(r"```(?:json)?\s*", "", text).strip()
     try:
         obj = json.loads(text)
-        if isinstance(obj, dict) and ("reasoning" in obj or "answer" in obj):
+        if isinstance(obj, dict) and "reasoning" in obj:
             return {
                 "reasoning": str(obj.get("reasoning","")),
                 "answer": str(obj.get("answer","")),
@@ -138,21 +123,6 @@ def reward_fn(completions, prompts=None, **kwargs):
             parsed = parse_action(text)
         except:
             parsed = FAIL.copy()
-
-        prompt_key = prompts[0][:50] if prompts else "default"
-        _episode_counts[prompt_key] = _episode_counts.get(prompt_key, 0) + 1
-        if parsed.get("abstain", False):
-            _abstain_counts[prompt_key] = _abstain_counts.get(prompt_key, 0) + 1
-
-        abstain_rate = _abstain_counts.get(prompt_key, 0) / max(1, _episode_counts.get(prompt_key, 1))
-
-        if parsed.get("abstain", False):
-            if abstain_rate > 0.3:
-                rewards.append(-0.5)
-            else:
-                rewards.append(0.0)
-            continue
-
         try:
             with _env_lock:
                 obs = _env.reset()
@@ -198,8 +168,8 @@ train_dataset = Dataset.from_list([
     for q in questions
 ])
 
-print("Starting Level 1 training...")
-wandb.init(project=WANDB_PROJECT, name="1.5b-level1-improved")
+print("Starting training...")
+wandb.init(project=WANDB_PROJECT, name="1.5b-level1-v2")
 
 trainer = GRPOTrainer(
     model=model,
@@ -209,13 +179,13 @@ trainer = GRPOTrainer(
         output_dir="/tmp/deceit-1.5b",
         bf16=torch.cuda.is_available() and torch.cuda.is_bf16_supported(),
         fp16=False,
-        max_steps=1000,
+        max_steps=500,
         per_device_train_batch_size=4,
         num_generations=4,
-        learning_rate=2e-5,
-        warmup_steps=10,
+        learning_rate=1e-5,
+        warmup_steps=5,
         logging_steps=1,
-        save_steps=100,
+        save_steps=50,
         report_to="wandb",
         max_completion_length=256,
         remove_unused_columns=False,
@@ -224,7 +194,7 @@ trainer = GRPOTrainer(
 )
 trainer.train()
 wandb.finish()
-print("Level 1 done!")
+print("Training done!")
 
 # Save Level 1 checkpoint
 model.save_pretrained("/tmp/deceit-1.5b-l1")
@@ -247,37 +217,6 @@ with open(data_path_l2) as f:
 
 print(f"Loaded {len(questions_l2)} Level 2 questions")
 
-data_path_l1 = pathlib.Path(_de.__file__).parent / "data" / "level1.jsonl"
-questions_l1 = []
-with open(data_path_l1) as f:
-    for line in f:
-        line = line.strip()
-        if line:
-            questions_l1.append(json.loads(line))
-
-# Mix 70% L2 + 30% L1
-n_l2 = len(questions_l2)
-n_l1_sample = max(1, int(n_l2 * 0.3))
-l1_sample = random.sample(questions_l1, min(n_l1_sample, len(questions_l1)))
-
-mixed_questions = []
-for q in questions_l2:
-    mixed_questions.append({
-        "question": q["question"],
-        "answer": q.get("answer", ""),
-        "distractors": q.get("distractors", []),
-        "is_l2": True
-    })
-for q in l1_sample:
-    mixed_questions.append({
-        "question": q["question"],
-        "answer": q.get("answer", ""),
-        "distractors": [],
-        "is_l2": False
-    })
-random.shuffle(mixed_questions)
-print(f"Mixed dataset: {len(mixed_questions)} questions ({n_l2} L2 + {len(l1_sample)} L1)")
-
 def make_prompt_l2(q, distractors):
     context = "\n".join(distractors)
     msgs = [
@@ -287,14 +226,12 @@ def make_prompt_l2(q, distractors):
     return tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
 
 train_dataset_l2 = Dataset.from_list([
-    {"prompt": make_prompt_l2(q["question"], q.get("distractors", [])),
-     "question": q["question"]}
-    for q in mixed_questions
+    {"prompt": make_prompt_l2(q["question"], q.get("distractors", [])), "question": q["question"]}
+    for q in questions_l2
 ])
 
+# Update env to level 2
 _env_l2 = DeceitEnvironment(grader=_grader)
-_abstain_counts_l2 = {}
-_episode_counts_l2 = {}
 
 def reward_fn_l2(completions, prompts=None, **kwargs):
     rewards = []
@@ -303,21 +240,6 @@ def reward_fn_l2(completions, prompts=None, **kwargs):
             parsed = parse_action(text)
         except:
             parsed = FAIL.copy()
-
-        prompt_key = prompts[0][:50] if prompts else "default"
-        _episode_counts_l2[prompt_key] = _episode_counts_l2.get(prompt_key, 0) + 1
-        if parsed.get("abstain", False):
-            _abstain_counts_l2[prompt_key] = _abstain_counts_l2.get(prompt_key, 0) + 1
-
-        abstain_rate = _abstain_counts_l2.get(prompt_key, 0) / max(1, _episode_counts_l2.get(prompt_key, 1))
-
-        if parsed.get("abstain", False):
-            if abstain_rate > 0.3:
-                rewards.append(-0.5)
-            else:
-                rewards.append(0.0)
-            continue
-
         try:
             with _env_lock:
                 obs = _env_l2.reset(level=2)
@@ -343,8 +265,9 @@ def reward_fn_l2(completions, prompts=None, **kwargs):
         rewards.append(total)
     return rewards
 
+# Train Level 2
 print("Starting Level 2 training on 1.5B...")
-wandb.init(project=WANDB_PROJECT, name="1.5b-level2-improved")
+wandb.init(project=WANDB_PROJECT, name="1.5b-level2-v2")
 
 trainer_l2 = GRPOTrainer(
     model=model,
@@ -352,15 +275,13 @@ trainer_l2 = GRPOTrainer(
     reward_funcs=[reward_fn_l2],
     args=GRPOConfig(
         output_dir="/tmp/deceit-1.5b-l2",
-        bf16=torch.cuda.is_available() and torch.cuda.is_bf16_supported(),
-        fp16=False,
-        max_steps=600,
+        max_steps=300,
         per_device_train_batch_size=4,
         num_generations=4,
-        learning_rate=2e-5,
-        warmup_steps=10,
+        learning_rate=2e-6,
+        warmup_steps=5,
         logging_steps=1,
-        save_steps=100,
+        save_steps=40,
         report_to="wandb",
         max_completion_length=256,
         remove_unused_columns=False,