{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  },
  "accelerator": "GPU",
  "gpuClass": "standard"
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Ballot Guide Lab — 4B control comparison\n",
    "Same v2 cases and pipeline; a new 4B LoRA adapter. Use a free Colab GPU. No paid endpoint.\n",
    "The old 2B adapter cannot be loaded into 4B. This experiment trains from the same 104 labels and 31 validation examples. Scores and human review remain pending until real outputs are graded.\n",
    "\nThe first attempt lost its runtime before completion could be verified. This retry notebook adds persistent Drive storage before inference/training. Mounting gives this notebook broad Drive access; code writes only a new BallotGuide4B folder. Do not grant that access unless you approve it.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare isolated project and pinned environment"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "from pathlib import Path\n",
    "from urllib.request import urlopen\n",
    "import json, hashlib, zipfile, sys, subprocess, os\n",
    "ROOT = Path('/content/ballot-4b-comparison/ballot-guide-lab')\n",
    "PROJECT_URL = 'https://www.utkarshruhela.com/votereducation/files/prototype-v2-voter-questions.zip'\n",
    "PROJECT_SHA256 = '54d5a7754e13b7b19ff952ccd4c2b772094dfbe9b728af9ff2a6d77eb04a3533'\n",
    "CASE_SHA256 = '97ad117af35b5e56589435ab25cf7591da724c06db983d7374aef966bdc38180'\n",
    "archive = Path('/content/ballot-4b-project.zip')\n",
    "archive.write_bytes(urlopen(PROJECT_URL, timeout=90).read())\n",
    "assert hashlib.sha256(archive.read_bytes()).hexdigest() == PROJECT_SHA256, 'Project archive changed.'\n",
    "destination = ROOT.parent\n",
    "with zipfile.ZipFile(archive) as z:\n",
    "    for name in z.namelist():\n",
    "        assert (destination/name).resolve().is_relative_to(destination.resolve()), 'Unsafe archive path.'\n",
    "    z.extractall(destination)\n",
    "assert hashlib.sha256((ROOT/'pilot/cases.json').read_bytes()).hexdigest() == CASE_SHA256\n",
    "sys.path.insert(0,str(ROOT)); sys.path.insert(0,str(ROOT/'training'))\n",
    "os.environ['UNSLOTH_COMPILE_DISABLE']='1'\n",
    "os.environ['GRADIO_ANALYTICS_ENABLED']='False'\n",
    "subprocess.check_call([sys.executable,'-m','pip','install','--quiet','--upgrade','uv'])\n",
    "import numpy, PIL\n",
    "deps=['torch==2.8.0','triton==3.4.0',f'numpy=={numpy.__version__}',f'pillow=={PIL.__version__}',\n",
    "      'torchvision==0.23.0','torchaudio==2.8.0','bitsandbytes==0.50.2','xformers==0.0.32.post2',\n",
    "      'unsloth_zoo[base] @ git+https://github.com/unslothai/unsloth-zoo@f787f63d3411c00633d1768f617348ee6867ed37',\n",
    "      'unsloth[base] @ git+https://github.com/unslothai/unsloth@7d2bd0118422c93840ea524303de679a1dc916cd']\n",
    "subprocess.check_call(['uv','pip','install','--system','--quiet',*deps])\n",
    "subprocess.check_call(['uv','pip','install','--system','--quiet','--no-deps','torchcodec==0.7.0','trl==0.22.2','torchao==0.18.0'])\n",
    "subprocess.check_call(['uv','pip','install','--system','--quiet','transformers==5.2.0','pypdf==6.10.0',\n",
    "                       'tokenizers==0.22.2','datasets==4.8.5','peft==0.21.1','huggingface-hub==1.33.0'])\n",
    "subprocess.check_call(['unsloth','install-kernels'])\n",
    "pipeline=ROOT/'training/colab_pipeline.py'\n",
    "code=pipeline.read_text()\n",
    "assert \"MODEL_ID = 'Qwen/Qwen3.5-2B'\" in code\n",
    "code=code.replace(\"MODEL_ID = 'Qwen/Qwen3.5-2B'\", \"MODEL_ID = 'Qwen/Qwen3.5-4B'\")\n",
    "code=code.replace(\"MODEL_REVISION = '15852e8c16360a2fea060d615a32b45270f8a8fc'\", \"MODEL_REVISION = '851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a'\")\n",
    "code=code.replace('results/colab_lora_v2_voterquestions','results/colab_4b_v2_control')\n",
    "pipeline.write_text(code)\n",
    "reference_archive=Path('/content/ballot-2b-reference.zip')\n",
    "reference_archive.write_bytes(urlopen('https://www.utkarshruhela.com/votereducation/files/v2-generated-results.zip',timeout=90).read())\n",
    "assert hashlib.sha256(reference_archive.read_bytes()).hexdigest()=='57beb1edae9b8915de1127925a0a871eb8783afa279e469126f03e3d4a3acfcd'\n",
    "with zipfile.ZipFile(reference_archive) as z:\n",
    "    reference_rows=[json.loads(line) for line in z.read('results/colab_lora_v2_voterquestions/outputs.jsonl').decode().splitlines() if line.strip()]\n",
    "(ROOT/'reference_2b_outputs.jsonl').write_text(''.join(json.dumps(r,ensure_ascii=False)+'\\n' for r in reference_rows))\n",
    "print('COMPARISON_SETUP_COMPLETE',ROOT, '4B control; unchanged sources/questions/training labels.', flush=True)\n"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Persist outputs and epoch checkpoints before the long run\n",
    "This step grants Colab Drive access. Outputs and checkpoints are written through links into a new task folder; source-package historical artifacts are kept separately.\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Requires your explicit approval for Colab to mount Google Drive.\n",
    "from google.colab import drive\n",
    "drive.mount('/content/drive')\n",
    "from datetime import datetime,timezone\n",
    "BACKUP_ROOT=Path('/content/drive/MyDrive/BallotGuide4B')/datetime.now(timezone.utc).strftime('%Y%m%dT%H%M%SZ')\n",
    "BACKUP_ROOT.mkdir(parents=True,exist_ok=False)\n",
    "for name in ['results','artifacts']:\n",
    "    link=ROOT/name\n",
    "    if link.is_symlink():\n",
    "        raise RuntimeError('Existing backup links detected; preserve this run and use a fresh runtime.')\n",
    "    if link.exists():\n",
    "        archive=ROOT/('source_package_'+name)\n",
    "        assert not archive.exists()\n",
    "        link.rename(archive)\n",
    "    destination=BACKUP_ROOT/name\n",
    "    destination.mkdir()\n",
    "    link.symlink_to(destination,target_is_directory=True)\n",
    "print('DURABLE_BACKUP_FOLDER',BACKUP_ROOT)\n",
    "print('Answers and epoch checkpoints will persist in Drive; model downloads remain in the local cache.')\n"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load 4B and verify the control"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Release previous model memory when appended to an already used runtime.\n",
    "import gc,sys,os\n",
    "for name in ['model','tokenizer','backend','trainer','stats']:\n",
    "    if name in globals(): globals()[name]=None\n",
    "for name in ['colab_pipeline','engine','sources']:\n",
    "    sys.modules.pop(name,None)\n",
    "gc.collect()\n",
    "from unsloth import FastLanguageModel\n",
    "import torch,hashlib,json,importlib.metadata,time\n",
    "from pathlib import Path\n",
    "ROOT=Path('/content/ballot-4b-comparison/ballot-guide-lab')\n",
    "sys.path=[str(ROOT/'training'),str(ROOT)]+[p for p in sys.path if '/content/ballot-' not in p]\n",
    "assert torch.cuda.is_available(), 'Select a free GPU runtime before running.'\n",
    "torch.cuda.empty_cache();torch._dynamo.config.disable=True\n",
    "from colab_pipeline import MODEL_ID,MODEL_REVISION,MAX_LENGTH,load_base,train_lora,audit_data,generation_backend\n",
    "from sources import CONFIG,index_sources,packet_for\n",
    "assert MODEL_ID=='Qwen/Qwen3.5-4B'\n",
    "assert MODEL_REVISION=='851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a'\n",
    "data_manifest=audit_data()\n",
    "assert data_manifest['files']['train']['examples']==104 and data_manifest['files']['validation']['examples']==31\n",
    "records=index_sources();cases=json.loads((ROOT/'pilot/cases.json').read_text())\n",
    "assert len(cases)==15 and sum(c['difficult'] for c in cases)==5\n",
    "reference=[json.loads(l) for l in (ROOT/'reference_2b_outputs.jsonl').read_text().splitlines()]\n",
    "packets={c['id']:packet_for(c['proposition'],records,c['question']) for c in cases}\n",
    "for case in cases:\n",
    "    old=next(r for r in reference if r['case_id']==case['id'])\n",
    "    assert packets[case['id']]==old['retrieved_sources'], 'Retrieval changed; stop this control experiment.'\n",
    "model,tokenizer=load_base()\n",
    "assert not getattr(model,'peft_config',None), 'The untuned condition must not contain a trained adapter.'\n",
    "run_dir=ROOT/'results/colab_4b_v2_control';run_dir.mkdir(parents=True,exist_ok=True)\n",
    "manifest={'run_name':'colab_4b_v2_control','test_version':'v2-voterquestions',\n",
    " 'base_model':MODEL_ID,'base_revision':MODEL_REVISION,'gpu':torch.cuda.get_device_name(0),\n",
    " 'control':'Same v2 cases, prompts, source packets, synthetic labels, LoRA recipe and decoder as the 2B experiment; only base checkpoint changes.',\n",
    " 'source_packets_identical_to_2b':True,'source_pdf_sha256':CONFIG['source_sha256'],\n",
    " 'cases_sha256':hashlib.sha256((ROOT/'pilot/cases.json').read_bytes()).hexdigest(),\n",
    " 'prototype_prompt_sha256':hashlib.sha256((ROOT/'prompts/prototype.txt').read_bytes()).hexdigest(),\n",
    " 'baseline_prompt_sha256':hashlib.sha256((ROOT/'prompts/baseline.txt').read_bytes()).hexdigest(),\n",
    " 'decoder':'greedy Transformers; no JSON grammar','max_tokens':320,'max_length':3072,\n",
    " 'seed':894,'enable_thinking':False,'quantized':False,'training_required':True,\n",
    " 'packages':{n:importlib.metadata.version(n) for n in ['torch','transformers','unsloth','unsloth_zoo','peft','pypdf']},\n",
    " 'started_at':__import__('datetime').datetime.now(__import__('datetime').timezone.utc).isoformat()}\n",
    "assert not (run_dir/'outputs.jsonl').exists(), 'Use a fresh directory; do not overwrite existing answers.'\n",
    "(run_dir/'runtime.json').write_text(json.dumps(manifest,indent=2))\n",
    "print('MODEL_4B_READY',MODEL_REVISION,'GPU:',torch.cuda.get_device_name(0),'source packets match 2B.',flush=True)\n"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Freeze the 30 untuned answers"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "from datetime import datetime,timezone\n",
    "\n",
    "def generate_condition(label,conditions):\n",
    "    backend=generation_backend(model,tokenizer)\n",
    "    output=run_dir/'outputs.jsonl'\n",
    "    previous=[json.loads(l) for l in output.read_text().splitlines()] if output.exists() else []\n",
    "    done={(r['model_label'],r['condition'],r['case_id']) for r in previous}\n",
    "    for case in cases:\n",
    "        for condition in conditions:\n",
    "            key=(label,condition,case['id'])\n",
    "            if key in done:continue\n",
    "            torch.manual_seed(894)\n",
    "            result=backend.answer(case['question'],case['proposition'],condition,records)\n",
    "            old=next(r for r in reference if r['case_id']==case['id'] and r['condition']==condition)\n",
    "            assert result['prompt_sha256']==old['prompt_sha256'], 'Prompt changed from 2B control.'\n",
    "            result.update(model_label=label,case_id=case['id'],difficult=case['difficult'],\n",
    "                generated_at=datetime.now(timezone.utc).isoformat(),retrieved_sources=packets[case['id']],\n",
    "                decoder='greedy Transformers; no JSON grammar',test_version='v2-voterquestions',base_model=MODEL_ID)\n",
    "            with output.open('a') as f:f.write(json.dumps(result,ensure_ascii=False)+'\\n')\n",
    "            done.add(key)\n",
    "            print('COMPARISON_PROGRESS',len(done),'/45',label,condition,case['id'],result['status'],result['latency_seconds'],'s',flush=True)\n",
    "assert not getattr(model,'peft_config',None)\n",
    "generate_condition('untuned',['baseline','prototype'])\n",
    "print('UNTUNED_4B_COMPLETE: 30 actual outputs frozen before training.',flush=True)\n"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train a fresh 4B adapter using the original recipe"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "before_rows=[json.loads(l) for l in (run_dir/'outputs.jsonl').read_text().splitlines() if l.strip()]\n",
    "assert len(before_rows)==30 and all(r['model_label']=='untuned' for r in before_rows)\n",
    "before_sha256=hashlib.sha256((run_dir/'outputs.jsonl').read_bytes()).hexdigest()\n",
    "(run_dir/'untuned_freeze.json').write_text(json.dumps({'outputs':30,'sha256':before_sha256,'frozen_before_training':True},indent=2))\n",
    "torch.cuda.empty_cache()\n",
    "torch.cuda.reset_peak_memory_stats()\n",
    "model,trainer,training_manifest=train_lora(model,tokenizer)\n",
    "assert training_manifest['status']=='trained' and training_manifest['base_model']=='Qwen/Qwen3.5-4B'\n",
    "assert training_manifest['optimizer_steps']==52\n",
    "training_manifest['untuned_outputs_sha256']=before_sha256\n",
    "training_manifest['peak_gpu_allocated_bytes']=torch.cuda.max_memory_allocated()\n",
    "training_manifest['peak_gpu_reserved_bytes']=torch.cuda.max_memory_reserved()\n",
    "training_manifest['adapter_model_sha256']=hashlib.sha256((ROOT/'artifacts/lora_adapter/adapter_model.safetensors').read_bytes()).hexdigest()\n",
    "(ROOT/'artifacts/training_manifest.json').write_text(json.dumps(training_manifest,indent=2))\n",
    "print('TRAINED_4B_COMPLETE',json.dumps({'steps':training_manifest['optimizer_steps'],\n",
    " 'seconds':training_manifest['training_seconds'],'validation_loss':training_manifest['validation_metrics'].get('eval_loss'),\n",
    " 'adapter_sha256':training_manifest['adapter_model_sha256']}),flush=True)\n"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Freeze the 15 tuned answers and export evidence"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "assert training_manifest['status']=='trained' and training_manifest['base_model']=='Qwen/Qwen3.5-4B'\n",
    "assert training_manifest['optimizer_steps']==52\n",
    "assert hashlib.sha256((ROOT/'artifacts/lora_adapter/adapter_model.safetensors').read_bytes()).hexdigest()==training_manifest['adapter_model_sha256']\n",
    "FastLanguageModel.for_inference(model)\n",
    "assert all(layer.enabled for layer in model.get_layer_status())\n",
    "generate_condition('finetuned',['prototype'])\n",
    "rows=[json.loads(l) for l in (run_dir/'outputs.jsonl').read_text().splitlines() if l.strip()]\n",
    "assert len(rows)==45 and len({(r['model_label'],r['condition'],r['case_id']) for r in rows})==45\n",
    "untuned_bytes=b'\\n'.join((run_dir/'outputs.jsonl').read_bytes().splitlines()[:30])+b'\\n'\n",
    "assert hashlib.sha256(untuned_bytes).hexdigest()==training_manifest['untuned_outputs_sha256'], 'Untuned answers changed after training.'\n",
    "for case in cases:\n",
    "    group=[r for r in rows if r['case_id']==case['id']]\n",
    "    assert len(group)==3 and all(r['retrieved_sources']==packets[case['id']] for r in group)\n",
    "    paired=[r for r in group if r['condition']=='prototype']\n",
    "    assert paired[0]['prompt_sha256']==paired[1]['prompt_sha256']\n",
    "summary={'outputs':45,'case_count':15,'source_packets_and_prompts_match_2b':True,\n",
    " 'scores_complete':False,'human_ratings_complete':False,\n",
    " 'groups':{label+'_'+condition:{'outputs':len(g),'technical_passes':sum(not r['validation_errors'] for r in g) if condition=='prototype' else None} for label,condition in [('untuned','baseline'),('untuned','prototype'),('finetuned','prototype')] for g in [[r for r in rows if r['model_label']==label and r['condition']==condition]]}}\n",
    "(run_dir/'technical_summary.json').write_text(json.dumps(summary,indent=2))\n",
    "# Fresh shuffled IDs prevent the old 2B condition key from revealing 4B labels.\n",
    "evaluator=ROOT/'evaluate.py'\n",
    "evaluation_code=evaluator.read_text()\n",
    "assert 'random.Random(894).shuffle(shuffled)' in evaluation_code\n",
    "evaluator.write_text(evaluation_code.replace('random.Random(894).shuffle(shuffled)','random.Random(895).shuffle(shuffled)'))\n",
    "manifest['blinding_seed']=895\n",
    "subprocess.check_call([sys.executable,str(evaluator),'prepare','--run-dir',str(run_dir)])\n",
    "(run_dir/'test_cases.json').write_text(json.dumps(cases,indent=2))\n",
    "manifest['parameter_elements_by_dtype']={str(d):sum(p.numel() for p in model.parameters() if p.dtype==d) for d in {p.dtype for p in model.parameters()}}\n",
    "manifest['completed_at']=datetime.now(timezone.utc).isoformat()\n",
    "manifest['training_manifest']=training_manifest\n",
    "manifest['outputs_sha256']=hashlib.sha256((run_dir/'outputs.jsonl').read_bytes()).hexdigest()\n",
    "(run_dir/'runtime.json').write_text(json.dumps(manifest,indent=2))\n",
    "package=Path('/content/ballot-guide-4b-v2-results.zip')\n",
    "with zipfile.ZipFile(package,'w',zipfile.ZIP_DEFLATED) as z:\n",
    "    for p in run_dir.rglob('*'):\n",
    "        if p.is_file():z.write(p,p.relative_to(ROOT))\n",
    "    for rel in ['pilot/cases.json','pilot/rubric.json','prompts/prototype.txt','prompts/baseline.txt','prompts/judge.txt','training/colab_pipeline.py','evaluate.py','training/data_manifest.json','artifacts/training_manifest.json','artifacts/trainer_history.json','config.json','data/source_manifest.json']:\n",
    "        z.write(ROOT/rel,rel)\n",
    "print('COMPARISON_4B_GENERATION_COMPLETE',json.dumps(summary),flush=True)\n",
    "print('RESULTS_ARCHIVE_SHA256',hashlib.sha256(package.read_bytes()).hexdigest(),flush=True)\n",
    "from IPython.display import display,HTML\n",
    "import base64\n",
    "# Visible browser export; no hidden notebook or kernel APIs are needed.\n",
    "display(HTML('<textarea id=\"comparison-export\" style=\"width:100%;height:160px\">'+base64.b64encode(package.read_bytes()).decode()+'</textarea>'))\n"
   ],
   "execution_count": null,
   "outputs": []
  }
 ]
}
