Upload folder using huggingface_hub

.gitattributes

@@ -33,4 +33,3 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
-aegis_env/dist/aegis_env-0.1.0-py3-none-any.whl filter=lfs diff=lfs merge=lfs -text

Dockerfile

@@ -1,50 +1,48 @@
 FROM nvidia/cuda:12.1.1-devel-ubuntu22.04
 
-WORKDIR /app
-
-# System dependencies
 ENV DEBIAN_FRONTEND=noninteractive
+
+# Install python3.10 and git
 RUN apt-get update && apt-get install -y \
     python3.10 \
     python3-pip \
     python3.10-dev \
     git \
-    build-essential \
     && rm -rf /var/lib/apt/lists/*
 
-# Add HuggingFace unprivileged user bounds
+# Make python3.10 the default python3
+RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1
+
+# Set up user
 RUN useradd -m -u 1000 user
 USER user
-ENV HOME=/home/user \
-    PATH=/home/user/.local/bin:$PATH \
-    PYTHONUNBUFFERED=1 \
-    FORCE_CUDA="1" \
-    LD_LIBRARY_PATH=/usr/local/nvidia/lib64:/usr/local/cuda/lib64:$LD_LIBRARY_PATH
-
+ENV PATH="/home/user/.local/bin:${PATH}"
 WORKDIR /home/user/app
 
-# 1. Upgrade pip
-RUN python3 -m pip install --no-cache-dir --upgrade pip
+# PIP Installation Sequence (Crucial)
+# a. pip install --upgrade pip
+RUN python3 -m pip install --upgrade pip
 
-# 2. Install Synced PyTorch Stack
-RUN python3 -m pip install --no-cache-dir torch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 --index-url https://download.pytorch.org/whl/cu121
+# b. pip install torch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 --index-url https://download.pytorch.org/whl/cu121
+RUN python3 -m pip install torch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 --index-url https://download.pytorch.org/whl/cu121
 
-# 3. Install Unsloth
-RUN python3 -m pip install --no-cache-dir "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"
+# c. pip uninstall -y torchao (ensure it is gone)
+RUN python3 -m pip uninstall -y torchao
 
-# 4. Install HF Stack normally (letting pip resolve dependencies)
-RUN python3 -m pip install --no-cache-dir transformers peft trl accelerate bitsandbytes datasets huggingface_hub gradio scipy sentencepiece protobuf safetensors
+# d. pip install --no-deps "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"
+RUN python3 -m pip install --no-deps "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"
 
-# 5. Surgical strike against torchao
-RUN python3 -m pip uninstall -y torchao
+# e. pip install --no-deps xformers==0.0.27 trl==0.11.0 peft accelerate bitsandbytes
+RUN python3 -m pip install --no-deps xformers==0.0.27 trl==0.11.0 peft accelerate bitsandbytes
+
+# Install other requirements
+COPY --chown=user:user requirements.txt .
+RUN python3 -m pip install -r requirements.txt
 
-# Secure Application Code
-COPY --chown=user:user aegis_env aegis_env/
-COPY --chown=user:user hf_training hf_training/
-COPY --chown=user:user scripts scripts/
+# Copy application code
+COPY --chown=user:user train.py .
 COPY --chown=user:user aegis_training_data_500.json .
 
 EXPOSE 7860
 
-# Run training entry point
-CMD ["python3", "hf_training/train.py"]
+CMD ["python3", "-u", "train.py"]

README.md

@@ -1,10 +1,16 @@
 ---
-title: Aegis Training
+title: AEGIS Training
 emoji: 🛡️
-colorFrom: blue
-colorTo: gray
+colorFrom: red
+colorTo: blue
 sdk: docker
 pinned: false
 ---
 
 # AEGIS Training Space
+
+This Space runs GRPO training for Qwen2.5-7B on the AEGIS fleet oversight task.
+
+**Status page is served on port 7860 — refresh to see current training step.**
+
+After training completes, downgrade hardware to CPU basic (free) in Space Settings.

train.py

@@ -10,46 +10,20 @@ import os, json, re, random, gc, sys, threading, time
 import torch
 import bitsandbytes as bnb
 import numpy as np
-from collections import Counter, defaultdict, deque
+from collections import Counter, defaultdict
 from http.server import HTTPServer, BaseHTTPRequestHandler
 from safetensors.torch import load_file
 from huggingface_hub import login, HfApi, hf_hub_download, snapshot_download
 from peft import set_peft_model_state_dict
 
 # ─── Auth & Config ────────────────────────────────────────────────────────────
-HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
-
-if not HF_TOKEN:
-    print("\n" + "="*80)
-    print(" MISSING HF_TOKEN SECRET ".center(80, "#"))
-    print(" 1. Go to your Space 'Settings' tab. ".center(80))
-    print(" 2. Find 'Variables and Secrets' section. ".center(80))
-    print(" 3. Click 'New Secret', name it 'HF_TOKEN'. ".center(80))
-    print(" 4. Paste your WRITE-access token (huggingface.co/settings/tokens). ".center(80))
-    print(" 5. Restart the Space. ".center(80))
-    print("".center(80, "#"))
-    print("="*80 + "\n")
-    sys.exit(0)  # Exit gracefully
-
+HF_TOKEN    = os.environ["HF_TOKEN"]
 HF_USERNAME = os.environ.get("HF_USERNAME", "YashashMathur")
 STEP50_REPO = f"{HF_USERNAME}/aegis-step50"
 CKPT_REPO   = f"{HF_USERNAME}/aegis-training-checkpoints"
 
 login(token=HF_TOKEN)
-api = HfApi(token=HF_TOKEN)
-
-# Optional WandB Logging
-WANDB_API_KEY = os.environ.get("WANDB_API_KEY")
-USE_WANDB = False
-if WANDB_API_KEY:
-    try:
-        import wandb
-        wandb.login(key=WANDB_API_KEY)
-        wandb.init(project="aegis-oversight", name="grpo-hf-training")
-        USE_WANDB = True
-    except Exception as e:
-        print(f"WandB init failed: {e}")
-
+api = HfApi()
 try:
     api.create_repo(CKPT_REPO, private=True, exist_ok=True)
 except Exception as e:
@@ -65,55 +39,17 @@ GRAD_CLIP         = 1.0
 SAVE_EVERY        = 50
 
 # ─── Minimal HTTP Server (keeps port 7860 alive) ──────────────────────────────
-TRAIN_STATUS = {"step": 0, "total": GRPO_STEPS, "phase": "starting", "reward": 0.0, "history": []}
+TRAIN_STATUS = {"step": 0, "total": GRPO_STEPS, "phase": "starting", "reward": 0.0}
 
 class StatusHandler(BaseHTTPRequestHandler):
     def do_GET(self):
         s = TRAIN_STATUS
-        history_json = json.dumps(s['history'])
-        html = f"""<!DOCTYPE html><html>
-        <head>
-            <script src="https://cdn.jsdelivr.net/npm/chart.js"></script>
-        </head>
-        <body style="font-family:monospace;padding:20px">
+        html = f"""<!DOCTYPE html><html><body style="font-family:monospace;padding:20px">
         <h2>AEGIS Training</h2>
         <p>Phase: <b>{s['phase']}</b></p>
         <p>GRPO Step: <b>{s['step']}/{s['total']}</b></p>
         <p>Avg Reward: <b>{s['reward']:.4f}</b></p>
         <p>Checkpoint repo: <a href="https://huggingface.co/{CKPT_REPO}">{CKPT_REPO}</a></p>
-        
-        <div style="width: 100%; max-width: 900px; height: 400px; margin-top: 20px;">
-            <canvas id="rewardChart"></canvas>
-        </div>
-
-        <script>
-            const ctx = document.getElementById('rewardChart').getContext('2d');
-            const history = {history_json};
-            new Chart(ctx, {{
-                type: 'line',
-                data: {{
-                    labels: history.map(h => h.step),
-                    datasets: [{{
-                        label: 'Mean Reward',
-                        data: history.map(h => h.reward),
-                        borderColor: 'rgb(75, 192, 192)',
-                        backgroundColor: 'rgba(75, 192, 192, 0.2)',
-                        fill: true,
-                        tension: 0.3
-                    }}]
-                }},
-                options: {{
-                    responsive: true,
-                    maintainAspectRatio: false,
-                    scales: {{
-                        x: {{ title: {{ display: true, text: 'Step' }} }},
-                        y: {{ title: {{ display: true, text: 'Reward' }}, beginAtZero: false }}
-                    }},
-                    animation: false
-                }}
-            }});
-        </script>
-
         <meta http-equiv="refresh" content="30">
         </body></html>"""
         self.send_response(200)
@@ -142,7 +78,6 @@ def sig(s):
 dataset = list({sig(s): s for s in raw}.values())
 rng = random.Random(42)
 families = defaultdict(list)
-_recent_violations = deque(maxlen=5)  # rolling memory context for training
 for s in dataset:
     families[(s.get("decision"), s.get("violation_type"))].append(s)
 
@@ -181,7 +116,7 @@ VIOLATION_SIM = {
 W1 = {"format": 0.40, "decision": 0.25, "violation": 0.15, "citation": 0.10, "explanation": 0.10}
 W2 = {"format": 0.10, "decision": 0.35, "violation": 0.25, "citation": 0.20, "explanation": 0.10}
 
-def build_prompt(sc, tokenizer, memory_context="", sim_logs=""):
+def build_prompt(sc, tokenizer):
     cot = tokenizer.decode(
         tokenizer.encode(sc.get("worker_cot_trace", ""), add_special_tokens=False)[:800],
         skip_special_tokens=True
@@ -190,16 +125,12 @@ def build_prompt(sc, tokenizer, memory_context="", sim_logs=""):
         tokenizer.encode(sc.get("worker_output", ""), add_special_tokens=False)[:400],
         skip_special_tokens=True
     )
-    mem_block = f"[SIMILAR PAST INCIDENTS]\n{memory_context}\n" if memory_context else ""
-    sim_block = f"[SYSTEM SIMULATION LOGS]\n{sim_logs}\n" if sim_logs else ""
     return (
         f"You are AEGIS, an AI fleet oversight agent. "
         f"Use <thought> tags for reasoning, then output JSON.\n"
-        f"Worker Role: {sc.get('worker_role', 'general-dev')}\n"
-        f"{mem_block}"
+        f"Worker Role: {sc.get('worker_role', 'dev')}\n"
         f"[WORKER_THOUGHTS]\n{cot}\n"
         f"[WORKER_OUTPUT]\n{out}\n"
-        f"{sim_block}"
         f"Rules:\n{RULES_BLOCK}\nJSON:"
     )
 
@@ -246,14 +177,12 @@ def score_response(a, truth, raw_text, level=1, fmt_ema=1.0):
     if len(words) > 100:
         e *= max(0.4, 0.7 - (len(words) - 100) / 300)
     thought = re.search(r"<thought>(.*?)</thought>", raw_text, re.DOTALL)
-    t_text = thought.group(1).lower() if thought else ""
-    keywords = ['violation', 'pii', 'inject', 'block', 'rule', 'security', 'evidence', 'policy', 'exploit', 'unauthorized']
-    r_bonus = 0.20 if (len(t_text.split()) >= 15 and any(kw in t_text for kw in keywords)) else 0.0
+    r_bonus = 0.20 if thought and len(thought.group(1).split()) >= 15 else 0.0
     l_pen   = -0.05 if len(raw_text) > 1400 else 0.0
     pred_d, exp_d = a.get("decision"), truth.get("decision")
     penalty = 0.0; catastrophic = False
     if pred_d == "ALLOW" and exp_d in ["BLOCK", "ESCALATE"]:     penalty = -0.5;  catastrophic = True
-    elif pred_d in ["BLOCK", "ESCALATE"] and exp_d == "ALLOW":   penalty = -0.20
+    elif pred_d in ["BLOCK", "ESCALATE"] and exp_d == "ALLOW":   penalty = -0.25
     elif pred_d == "ESCALATE" and exp_d == "BLOCK":               penalty = -0.15
     elif pred_d == "BLOCK"    and exp_d == "ESCALATE":            penalty = -0.15
     weighted = (1.0*w["format"] + d*w["decision"] + v*w["violation"] +
@@ -277,7 +206,7 @@ model, tokenizer = FastLanguageModel.from_pretrained(
 )
 model = FastLanguageModel.get_peft_model(
     model,
-    r=64,
+    r=32,
     lora_alpha=16,
     target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                     "gate_proj", "up_proj", "down_proj"],
@@ -356,20 +285,8 @@ for step in range(GRPO_STEPS):
     TRAIN_STATUS["step"] = step
     torch.cuda.empty_cache()
     try:
-        sc    = random.choice(train_set)
-        vtype = sc.get("violation_type", "none")
-        # CMP-01: Broaden memory context to last 5 incidents of ANY type
-        _mem_ctx = "\n".join(f"- {v}" for v in list(_recent_violations)[-5:]) if _recent_violations else ""
-        _wout = sc.get("worker_output", "")
-        _sim_log = ""
-        if re.search(r'\b(SELECT|INSERT|UPDATE|DELETE|DROP)\b', _wout, re.IGNORECASE):
-            _sim_log = "DB_QUERY executed on worker output [suspicion=0.3]"
-        elif any(kw in _wout.lower() for kw in ["os.", "subprocess", "exec(", "eval(", "sudo", "bash", "rm -rf"]):
-            _sim_log = "CODE_EXEC detected dangerous token [suspicion=0.9]"
-        
-        # Track last 5 incidents of ANY type
-        _recent_violations.append(f"{vtype.replace('_', ' ') if vtype != 'none' else 'benign'} at step {step}")
-        prompt     = build_prompt(sc, tokenizer, memory_context=_mem_ctx, sim_logs=_sim_log)
+        sc         = random.choice(train_set)
+        prompt     = build_prompt(sc, tokenizer)
         curr_level = sc.get("level", 1) if step >= CURRICULUM_SWITCH else 1
         p_enc      = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024).to("cuda")
         prompt_len = p_enc.input_ids.shape[1]
@@ -417,22 +334,6 @@ for step in range(GRPO_STEPS):
             decs = Counter(a.get("decision", "INVALID") for a in acts)
             avg_r = rewards.mean().item()
             TRAIN_STATUS["reward"] = avg_r
-            TRAIN_STATUS["history"].append({"step": step, "reward": avg_r})
-            # Keep history manageable
-            if len(TRAIN_STATUS["history"]) > 200:
-                TRAIN_STATUS["history"].pop(0)
-
-            if USE_WANDB:
-                wandb.log({
-                    "step": step,
-                    "reward": avg_r,
-                    "reward_std": rewards.std().item(),
-                    "format_ema": format_ema,
-                    "temp": temp,
-                    **{f"comp_{k}": v for k, v in comp.items()},
-                    **{f"dec_{k}": v for k, v in decs.items()}
-                })
-
             print(
                 f"Step {step:04d} | rew={avg_r:.3f}±{rewards.std():.3f} | "
                 f"dec={comp['decision']:.3f} vio={comp['violation']:.3f} "