feat: implement Unsloth training pipeline with Modal volumes and update documentation on model persona mitigation

README.md

@@ -65,3 +65,10 @@ modal run modal_train.py
 ```
 
 `modal_train.py` is currently a scaffold for the hackathon story. The remaining production step is to wire `verbs.jsonl` into Modal storage, uncomment the `SFTTrainer` block, run the LoRA job, and export/publish the GGUF or adapter artifact.
+
+## Known Limitations & Model Trade-offs
+
+During semantic stress testing of the live backend, we discovered the classic "Alignment Tax" associated with small, 4B-parameter local models:
+* **The Problem:** Highly instruction-tuned base models (like Nemotron-Mini and MiniCPM) are deeply trained to act like "helpful conversational assistants." When prompted with complex goals like *"write a python script"*, they often break out of the Cognitive Pacemaker persona and output helpful paragraphs instead of strict 8-word physical atomic actions.
+* **Current Mitigation:** We employ aggressive few-shot prompting, manual Jinja template overrides, and Python-side text post-processing to brutally truncate their output.
+* **The Ultimate Fix (Roadmap):** Instead of changing to a larger, slower 8B+ model, we plan to fully execute the `modal_train.py` Unsloth fine-tuning pipeline. By training a LoRA adapter specifically on `verbs.jsonl`, we will permanently "lobotomize" the chatbot persona, forcing the model to strictly output imperative verbs natively.

modal_train.py

@@ -12,16 +12,24 @@ image = (
         "accelerate",
         "bitsandbytes"
     )
+    .add_local_file("verbs.jsonl", "/root/verbs.jsonl")
 )
 
 app = modal.App("step-zero-finetune")
+vol = modal.Volume.from_name("step-zero-volume", create_if_missing=True)
 
-@app.function(image=image, gpu="A10G", timeout=1800)
+@app.function(
+    image=image, 
+    gpu="A10G", 
+    timeout=1800,
+    volumes={"/vol": vol}
+)
 def train_model():
-    from unsloth import FastLanguageModel
+    from unsloth import FastLanguageModel, is_bfloat16_supported
     from datasets import load_dataset
     from trl import SFTTrainer
     from transformers import TrainingArguments
+    import json
 
     print("Loading Base Model: Nemotron-Mini-4B...")
     max_seq_length = 1024
@@ -30,7 +38,7 @@ def train_model():
         model_name = "nvidia/Nemotron-Mini-4B-Instruct",
         max_seq_length = max_seq_length,
         dtype = None,
-        load_in_4bit = True, # Crucial for speed and fitting on consumer GPUs later
+        load_in_4bit = True,
     )
 
     model = FastLanguageModel.get_peft_model(
@@ -46,21 +54,65 @@ def train_model():
         use_rslora = False,
     )
 
-    # Note: Ensure you uploaded 'verbs.jsonl' to a Modal Volume, or fetch it via URL in real implementation
-    # dataset = load_dataset("json", data_files="verbs.jsonl", split="train")
-    print("Simulating dataset load...")
+    print("Loading and formatting verbs.jsonl dataset...")
+    dataset = load_dataset("json", data_files="/root/verbs.jsonl", split="train")
     
+    def format_prompt(examples):
+        texts = []
+        for instruction, output in zip(examples["instruction"], examples["output"]):
+            try:
+                task = json.loads(output)["task"]
+            except:
+                task = output
+                
+            parts = instruction.split(". Failures: ")
+            goal = parts[0].replace("Goal: ", "").strip()
+            failures = parts[1].split(".")[0].strip()
+            
+            system_msg = "You are a cognitive pacemaker. Break down goals into extremely tiny, atomic physical actions under 8 words."
+            prompt = f"<extra_id_0>System\n{system_msg}\n\n"
+            prompt += f"<extra_id_1>User\nGoal: {goal}\nCompleted Tasks: None\nFailures: {failures}\nOutput the NEXT step.\n"
+            prompt += f"<extra_id_1>Assistant\n{task}\n"
+            texts.append(prompt)
+        return {"text": texts}
+
+    dataset = dataset.map(format_prompt, batched=True)
+
     print("Starting LoRA fine-tuning for Syntactic Enforcement...")
-    # trainer = SFTTrainer( ... ) 
-    # trainer.train()
+    trainer = SFTTrainer(
+        model = model,
+        tokenizer = tokenizer,
+        train_dataset = dataset,
+        dataset_text_field = "text",
+        max_seq_length = max_seq_length,
+        dataset_num_proc = 2,
+        packing = False,
+        args = TrainingArguments(
+            per_device_train_batch_size = 2,
+            gradient_accumulation_steps = 4,
+            warmup_steps = 5,
+            max_steps = 60,
+            learning_rate = 2e-4,
+            fp16 = not is_bfloat16_supported(),
+            bf16 = is_bfloat16_supported(),
+            logging_steps = 1,
+            optim = "adamw_8bit",
+            weight_decay = 0.01,
+            lr_scheduler_type = "linear",
+            seed = 3407,
+            output_dir = "outputs",
+        ),
+    )
+    trainer.train()
     
     print("Exporting to GGUF...")
-    # model.save_pretrained_gguf("step-zero-nemotron", tokenizer, quantization_method = "q4_k_m")
+    model.save_pretrained_gguf("/vol/step-zero-nemotron", tokenizer, quantization_method = "q4_k_m")
     
-    print("Training complete. GGUF artifact generated.")
+    print("Training complete. GGUF artifact generated in /vol.")
 
 @app.local_entrypoint()
 def main():
     print("Submitting training job to Modal infrastructure...")
     train_model.remote()
-    print("Done! Download your weights from your Modal storage.")
+    print("Done! To download your weights, run:")
+    print("modal volume get step-zero-volume step-zero-nemotron-unsloth-Q4_K_M.gguf models/")

test_modal_mount.py

@@ -0,0 +1,2 @@
+import modal
+print(hasattr(modal.Image, 'add_local_file'))