notebooks/train_colab.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# 🌾 Sahel-Agri Voice AI — Fine-tune on Farmer Feedback\n",
    "\n",
    "**Run after collecting ≥10 corrections in the Space.**  \n",
    "First run? Use `bootstrap_repos.ipynb` instead to train the v0 Waxal adapter.\n",
    "\n",
    "This notebook fine-tunes the existing LoRA adapter using:\n",
    "- **Waxal baseline** (up to 500 samples) — keeps the model grounded\n",
    "- **Farmer corrections** (3× upsampled) — targeted improvement from real field use\n",
    "\n",
    "**Before running:** Runtime → Change runtime type → **T4 GPU**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 1 — GPU check\n",
    "import subprocess\n",
    "result = subprocess.run(['nvidia-smi'], capture_output=True, text=True)\n",
    "if result.returncode != 0:\n",
    "    raise RuntimeError('No GPU! Runtime → Change runtime type → T4 GPU')\n",
    "print(result.stdout[:500])\n",
    "print('✅ GPU ready')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 2 — Install dependencies (matching Space versions)\n",
    "!pip install -q \\\n",
    "    torch==2.11.0 torchaudio==2.11.0 \\\n",
    "    transformers==5.5.0 datasets==4.8.4 \\\n",
    "    accelerate==1.13.0 evaluate==0.4.2 \\\n",
    "    huggingface-hub==1.9.0 peft==0.18.1 \\\n",
    "    librosa==0.10.2 soundfile==0.12.1 \\\n",
    "    jiwer==3.0.4 pyyaml==6.0.2\n",
    "print('✅ Packages installed')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 3 — HuggingFace login\n",
    "# Colab: 🔑 icon (left sidebar) → Add new secret → name=HF_TOKEN\n",
    "# Kaggle: Add Data → add as Kaggle secret named HF_TOKEN\n",
    "import os\n",
    "try:\n",
    "    from google.colab import userdata  # type: ignore\n",
    "    HF_TOKEN = userdata.get('HF_TOKEN')\n",
    "except Exception:\n",
    "    HF_TOKEN = os.environ.get('HF_TOKEN', '')\n",
    "\n",
    "if not HF_TOKEN:\n",
    "    raise ValueError('HF_TOKEN not found — see instructions above.')\n",
    "\n",
    "from huggingface_hub import login\n",
    "login(token=HF_TOKEN, add_to_git_credential=False)\n",
    "\n",
    "SPACE_REPO_ID    = 'ous-sow/sahel-agri-voice'\n",
    "FEEDBACK_REPO_ID = 'ous-sow/sahel-agri-feedback'\n",
    "ADAPTER_REPO_ID  = 'ous-sow/sahel-agri-adapters'\n",
    "# Must match what the Space uses — whisper-small for cpu-basic, whisper-large-v3-turbo for GPU.\n",
    "WHISPER_MODEL_ID = 'openai/whisper-small'\n",
    "TRAIN_LANG       = 'bam'   # ← change to 'ful' for Fula\n",
    "\n",
    "print(f'✅ Logged in | training language: {TRAIN_LANG}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 4 — Download Space code and feedback corrections\n",
    "import json, shutil, sys\n",
    "from pathlib import Path\n",
    "from huggingface_hub import snapshot_download, hf_hub_download\n",
    "\n",
    "# Get Space code (contains src/, configs/)\n",
    "space_dir = Path(snapshot_download(\n",
    "    repo_id=SPACE_REPO_ID, repo_type='space', token=HF_TOKEN\n",
    "))\n",
    "sys.path.insert(0, str(space_dir))\n",
    "print(f'Space code: {space_dir}')\n",
    "\n",
    "# Download feedback corrections.jsonl\n",
    "jsonl_path = hf_hub_download(\n",
    "    repo_id=FEEDBACK_REPO_ID,\n",
    "    filename='corrections.jsonl',\n",
    "    repo_type='dataset',\n",
    "    token=HF_TOKEN,\n",
    ")\n",
    "with open(jsonl_path, encoding='utf-8') as f:\n",
    "    all_records = [json.loads(l) for l in f if l.strip()]\n",
    "\n",
    "corrections = [\n",
    "    r for r in all_records\n",
    "    if r.get('is_correction') and r['language'] == TRAIN_LANG\n",
    "]\n",
    "print(f'Total feedback records : {len(all_records)}')\n",
    "print(f'Corrections for {TRAIN_LANG}  : {len(corrections)}')\n",
    "\n",
    "if len(corrections) < 5:\n",
    "    print('⚠️  Very few corrections — consider collecting more before training.')\n",
    "    print('   Training will proceed with Waxal only (corrections will be skipped).')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 5 — Download feedback audio files from HF Dataset repo\n",
    "fb_audio_dir = Path('/tmp/sahel_feedback_audio')\n",
    "fb_audio_dir.mkdir(exist_ok=True)\n",
    "\n",
    "skipped = 0\n",
    "for rec in corrections:\n",
    "    local_path = fb_audio_dir / Path(rec['audio_file']).name\n",
    "    if local_path.exists():\n",
    "        continue\n",
    "    try:\n",
    "        dl = hf_hub_download(\n",
    "            repo_id=FEEDBACK_REPO_ID,\n",
    "            filename=rec['audio_file'],\n",
    "            repo_type='dataset',\n",
    "            token=HF_TOKEN,\n",
    "        )\n",
    "        shutil.copy(dl, local_path)\n",
    "    except Exception as e:\n",
    "        skipped += 1\n",
    "        print(f'  skip {rec[\"audio_file\"]}: {e}')\n",
    "\n",
    "# Point records at local paths\n",
    "for rec in corrections:\n",
    "    local = fb_audio_dir / Path(rec['audio_file']).name\n",
    "    if local.exists():\n",
    "        rec['audio_file'] = str(local)\n",
    "\n",
    "available = [r for r in corrections if Path(r['audio_file']).exists()]\n",
    "print(f'Downloaded {len(available)} / {len(corrections)} audio files (skipped {skipped})')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 6 — Fine-tune: Waxal baseline + farmer corrections\n",
    "#\n",
    "# WhisperLoRATrainer.setup() loads Waxal (streaming).\n",
    "# merge_extra_data() materialises Waxal (up to 500 samples),\n",
    "# appends corrections (3× upsampled), shuffles the combined dataset.\n",
    "# train() runs standard Seq2SeqTrainer on the merged dataset.\n",
    "\n",
    "import os\n",
    "os.environ['HF_TOKEN'] = HF_TOKEN\n",
    "\n",
    "from src.training.trainer import WhisperLoRATrainer\n",
    "\n",
    "lang_config_map = {'bam': 'lora_bambara.yaml', 'ful': 'lora_fula.yaml'}\n",
    "base_cfg  = str(space_dir / 'configs' / 'base_config.yaml')\n",
    "lang_cfg  = str(space_dir / 'configs' / lang_config_map[TRAIN_LANG])\n",
    "output_dir = f'/tmp/sahel_adapter_{TRAIN_LANG}'\n",
    "\n",
    "# Override output_dir so adapter saves to /tmp on Colab\n",
    "import yaml\n",
    "with open(lang_cfg) as f:\n",
    "    lang_config = yaml.safe_load(f)\n",
    "lang_config['output_dir'] = output_dir\n",
    "tmp_lang_cfg = f'/tmp/lora_{TRAIN_LANG}_tmp.yaml'\n",
    "with open(tmp_lang_cfg, 'w') as f:\n",
    "    yaml.dump(lang_config, f)\n",
    "\n",
    "trainer = WhisperLoRATrainer(\n",
    "    base_config_path=base_cfg,\n",
    "    language_config_path=tmp_lang_cfg,\n",
    ")\n",
    "trainer.setup()\n",
    "\n",
    "if available:\n",
    "    print(f'Merging {len(available)} corrections (×3) with Waxal baseline (cap=500)...')\n",
    "    trainer.merge_extra_data(available, repeat=3, waxal_cap=500)\n",
    "else:\n",
    "    print('No corrections available — training on Waxal only.')\n",
    "\n",
    "trainer.train()\n",
    "print(f'✅ Training complete — adapter at {output_dir}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 7 — Push adapter to HF Model repo\n",
    "from huggingface_hub import HfApi\n",
    "api = HfApi(token=HF_TOKEN)\n",
    "\n",
    "path_in_repo = 'adapters/bambara' if TRAIN_LANG == 'bam' else 'adapters/fula'\n",
    "n_corrections = len(available)\n",
    "\n",
    "api.upload_folder(\n",
    "    folder_path=output_dir,\n",
    "    repo_id=ADAPTER_REPO_ID,\n",
    "    repo_type='model',\n",
    "    path_in_repo=path_in_repo,\n",
    "    commit_message=(\n",
    "        f'Fine-tune {TRAIN_LANG}: Waxal baseline + {n_corrections} farmer corrections'\n",
    "    ),\n",
    ")\n",
    "print(f'✅ Pushed to {ADAPTER_REPO_ID}/{path_in_repo}')\n",
    "print('\\nNext: Space → Tab 3 → Reload Adapters from Hub')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 8 — Sanity check: compare WER before vs after adapter\n",
    "import random, torch, librosa, jiwer\n",
    "from transformers import WhisperForConditionalGeneration, WhisperProcessor\n",
    "from peft import PeftModel\n",
    "\n",
    "if not available:\n",
    "    print('No test samples — skipping sanity check.')\n",
    "else:\n",
    "    test_rec = random.choice(available)\n",
    "    print(f'Audio    : {Path(test_rec[\"audio_file\"]).name}')\n",
    "    print(f'Expected : {test_rec[\"corrected_text\"]}')\n",
    "    print(f'Pre-train: {test_rec[\"whisper_output\"]}')\n",
    "\n",
    "    # Load base + adapter\n",
    "    processor = WhisperProcessor.from_pretrained(WHISPER_MODEL_ID, token=HF_TOKEN)\n",
    "    base = WhisperForConditionalGeneration.from_pretrained(\n",
    "        WHISPER_MODEL_ID, torch_dtype=torch.float16, token=HF_TOKEN\n",
    "    ).to('cuda')\n",
    "    model = PeftModel.from_pretrained(base, output_dir).eval()\n",
    "\n",
    "    audio_np, _ = librosa.load(test_rec['audio_file'], sr=16000, mono=True)\n",
    "    feats = processor.feature_extractor(\n",
    "        audio_np, sampling_rate=16000, return_tensors='pt'\n",
    "    ).input_features.half().to('cuda')\n",
    "\n",
    "    with torch.no_grad():\n",
    "        ids = model.generate(feats, max_new_tokens=256)\n",
    "    result = processor.batch_decode(ids, skip_special_tokens=True)[0].strip()\n",
    "    print(f'Post-train: {result}')\n",
    "\n",
    "    ref = test_rec['corrected_text']\n",
    "    wer_before = jiwer.wer(ref, test_rec['whisper_output']) if test_rec.get('whisper_output') else 1.0\n",
    "    wer_after  = jiwer.wer(ref, result)\n",
    "    print(f'\\nWER before: {wer_before:.1%}  →  WER after: {wer_after:.1%}')\n",
    "    if wer_after < wer_before:\n",
    "        print('✅ Adapter improved transcription quality!')\n",
    "    else:\n",
    "        print('ℹ️  No improvement on this single sample — collect more corrections and retrain.')"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}