{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# 🌾 Sahel-Agri Voice AI — One-Time Bootstrap\n",
    "\n",
    "**Run this notebook ONCE** before deploying your Space. It:\n",
    "\n",
    "1. Creates the three HuggingFace repos (`sahel-agri-feedback`, `sahel-agri-adapters`, `sahel-agri-voice`)\n",
    "2. Seeds the feedback dataset with a `corrections.jsonl` placeholder\n",
    "3. Trains v0 LoRA adapters for **Bambara** and **Fula** on the full Google Waxal dataset\n",
    "4. Pushes adapters to `ous-sow/sahel-agri-adapters`\n",
    "\n",
    "After this notebook completes, push your project code to the Space and your app will start\n",
    "with working Bambara/Fula speech recognition from day 1 — **no user corrections needed yet**.\n",
    "\n",
    "For subsequent improvement runs (after collecting farmer feedback), use `train_colab.ipynb`.\n",
    "\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\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",
    "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(\n",
    "        'HF_TOKEN not found.\\n'\n",
    "        'Colab: click the 🔑 icon → Add new secret → name=HF_TOKEN'\n",
    "    )\n",
    "\n",
    "from huggingface_hub import login, HfApi\n",
    "login(token=HF_TOKEN, add_to_git_credential=False)\n",
    "api = HfApi(token=HF_TOKEN)\n",
    "\n",
    "HF_USERNAME      = 'ous-sow'\n",
    "FEEDBACK_REPO_ID = f'{HF_USERNAME}/sahel-agri-feedback'\n",
    "ADAPTER_REPO_ID  = f'{HF_USERNAME}/sahel-agri-adapters'\n",
    "SPACE_REPO_ID    = f'{HF_USERNAME}/sahel-agri-voice'\n",
    "# whisper-small trains on Colab T4 in ~25 min and runs on CPU in ~10s.\n",
    "# Change to 'openai/whisper-large-v3-turbo' only if you upgrade to a GPU Space.\n",
    "WHISPER_MODEL_ID = 'openai/whisper-small'\n",
    "\n",
    "print(f'✅ Logged in as {HF_USERNAME}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 4 — Create HuggingFace repos (skips if they already exist)\n",
    "from huggingface_hub import RepoUrl\n",
    "\n",
    "def create_repo_if_missing(repo_id, repo_type, private=True):\n",
    "    try:\n",
    "        url = api.create_repo(\n",
    "            repo_id=repo_id,\n",
    "            repo_type=repo_type,\n",
    "            private=private,\n",
    "            exist_ok=True,\n",
    "        )\n",
    "        print(f'  ✅ {repo_type}: {repo_id}')\n",
    "        return url\n",
    "    except Exception as e:\n",
    "        print(f'  ⚠️  {repo_id}: {e}')\n",
    "\n",
    "print('Creating repos...')\n",
    "create_repo_if_missing(FEEDBACK_REPO_ID, 'dataset',  private=True)\n",
    "create_repo_if_missing(ADAPTER_REPO_ID,  'model',    private=True)\n",
    "create_repo_if_missing(SPACE_REPO_ID,    'space',    private=False)\n",
    "\n",
    "# Seed the feedback dataset with an empty corrections.jsonl\n",
    "import io\n",
    "try:\n",
    "    api.upload_file(\n",
    "        path_or_fileobj=io.BytesIO(b''),\n",
    "        path_in_repo='corrections.jsonl',\n",
    "        repo_id=FEEDBACK_REPO_ID,\n",
    "        repo_type='dataset',\n",
    "        commit_message='Init: empty corrections.jsonl',\n",
    "    )\n",
    "    print(f'  ✅ {FEEDBACK_REPO_ID}/corrections.jsonl initialised')\n",
    "except Exception as e:\n",
    "    print(f'  ⚠️  corrections.jsonl upload: {e} (may already exist)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 5 — Clone Space code (so we can use src/ and configs/)\n",
    "# If the Space is brand new and has no code yet, clone from the local zip instead.\n",
    "import sys\n",
    "from pathlib import Path\n",
    "from huggingface_hub import snapshot_download\n",
    "\n",
    "try:\n",
    "    space_dir = Path(snapshot_download(\n",
    "        repo_id=SPACE_REPO_ID, repo_type='space', token=HF_TOKEN\n",
    "    ))\n",
    "    print(f'Space code: {space_dir}')\n",
    "except Exception as e:\n",
    "    print(f'Could not download Space ({e})')\n",
    "    print('Uploading project code to Space first...')\n",
    "    # If you have the project on Colab already (e.g. mounted Drive), set:\n",
    "    # space_dir = Path('/content/drive/MyDrive/voice-model')\n",
    "    # Otherwise upload via git (see README step 6) and re-run this cell.\n",
    "    raise RuntimeError(\n",
    "        'Push your project to the Space first:\\n'\n",
    "        '  git remote add space https://huggingface.co/spaces/ous-sow/sahel-agri-voice\\n'\n",
    "        '  git push space main\\n'\n",
    "        'Then re-run this notebook.'\n",
    "    )\n",
    "\n",
    "sys.path.insert(0, str(space_dir))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 6 — Train v0 Bambara adapter on full Waxal (bam)\n",
    "#\n",
    "# Uses streaming — Waxal is ~4h of audio, we cap at 2000 samples for Colab budget.\n",
    "# Full training (~4000 steps) on the entire dataset: use a Kaggle P100 (12h limit).\n",
    "import os, yaml\n",
    "os.environ['HF_TOKEN'] = HF_TOKEN\n",
    "\n",
    "from src.training.trainer import WhisperLoRATrainer\n",
    "\n",
    "WAXAL_CAP = 2000  # raise to 10000+ on Kaggle for a stronger v0 model\n",
    "\n",
    "base_cfg    = str(space_dir / 'configs' / 'base_config.yaml')\n",
    "bam_cfg_src = str(space_dir / 'configs' / 'lora_bambara.yaml')\n",
    "bam_out     = '/tmp/sahel_adapter_bam'\n",
    "\n",
    "# Override output_dir\n",
    "with open(bam_cfg_src) as f:\n",
    "    bam_config = yaml.safe_load(f)\n",
    "bam_config['output_dir'] = bam_out\n",
    "tmp_bam_cfg = '/tmp/lora_bam.yaml'\n",
    "with open(tmp_bam_cfg, 'w') as f:\n",
    "    yaml.dump(bam_config, f)\n",
    "\n",
    "# Also override max_steps in base config to match Waxal cap\n",
    "with open(base_cfg) as f:\n",
    "    base_config = yaml.safe_load(f)\n",
    "# ~2 steps per sample @ batch_size=4, gradient_acc=4\n",
    "base_config['training']['max_steps'] = max(500, WAXAL_CAP // 8)\n",
    "tmp_base_cfg = '/tmp/base_config.yaml'\n",
    "with open(tmp_base_cfg, 'w') as f:\n",
    "    yaml.dump(base_config, f)\n",
    "\n",
    "print(f'Training Bambara v0 adapter (Waxal cap={WAXAL_CAP}, max_steps={base_config[\"training\"][\"max_steps\"]})...')\n",
    "trainer_bam = WhisperLoRATrainer(\n",
    "    base_config_path=tmp_base_cfg,\n",
    "    language_config_path=tmp_bam_cfg,\n",
    ")\n",
    "trainer_bam.setup()\n",
    "\n",
    "# No feedback yet — materialise Waxal and train\n",
    "trainer_bam.merge_extra_data([], repeat=1, waxal_cap=WAXAL_CAP)\n",
    "\n",
    "trainer_bam.train()\n",
    "print(f'✅ Bambara v0 adapter saved to {bam_out}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 7 — Train v0 Fula adapter on full Waxal (ful)\n",
    "ful_cfg_src = str(space_dir / 'configs' / 'lora_fula.yaml')\n",
    "ful_out     = '/tmp/sahel_adapter_ful'\n",
    "\n",
    "with open(ful_cfg_src) as f:\n",
    "    ful_config = yaml.safe_load(f)\n",
    "ful_config['output_dir'] = ful_out\n",
    "tmp_ful_cfg = '/tmp/lora_ful.yaml'\n",
    "with open(tmp_ful_cfg, 'w') as f:\n",
    "    yaml.dump(ful_config, f)\n",
    "\n",
    "print(f'Training Fula v0 adapter (Waxal cap={WAXAL_CAP})...')\n",
    "trainer_ful = WhisperLoRATrainer(\n",
    "    base_config_path=tmp_base_cfg,\n",
    "    language_config_path=tmp_ful_cfg,\n",
    ")\n",
    "trainer_ful.setup()\n",
    "trainer_ful.merge_extra_data([], repeat=1, waxal_cap=WAXAL_CAP)\n",
    "trainer_ful.train()\n",
    "print(f'✅ Fula v0 adapter saved to {ful_out}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 8 — Push both adapters to HF Model repo\n",
    "from huggingface_hub import HfApi\n",
    "api = HfApi(token=HF_TOKEN)\n",
    "\n",
    "for lang, out_dir, path_in_repo in [\n",
    "    ('bam', bam_out, 'adapters/bambara'),\n",
    "    ('ful', ful_out, 'adapters/fula'),\n",
    "]:\n",
    "    api.upload_folder(\n",
    "        folder_path=out_dir,\n",
    "        repo_id=ADAPTER_REPO_ID,\n",
    "        repo_type='model',\n",
    "        path_in_repo=path_in_repo,\n",
    "        commit_message=f'v0 {lang} adapter trained on Waxal (cap={WAXAL_CAP} samples)',\n",
    "    )\n",
    "    print(f'✅ {lang} → {ADAPTER_REPO_ID}/{path_in_repo}')\n",
    "\n",
    "print()\n",
    "print('Bootstrap complete!')\n",
    "print()\n",
    "print('Next steps:')\n",
    "print('  1. Push your project code to the Space (git push space main)')\n",
    "print('  2. In Space Settings → Secrets, add HF_TOKEN, FEEDBACK_REPO_ID, ADAPTER_REPO_ID')\n",
    "print('  3. Space will build — your app at https://huggingface.co/spaces/ous-sow/sahel-agri-voice')\n",
    "print('  4. Tab 3 → Reload Adapters — Bambara + Fula adapters will be loaded')\n",
    "print('  5. Collect farmer corrections, then run train_colab.ipynb to keep improving')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 9 — Quick verification: list what was pushed to the adapter repo\n",
    "from huggingface_hub import list_repo_files\n",
    "\n",
    "files = sorted(list_repo_files(ADAPTER_REPO_ID, repo_type='model', token=HF_TOKEN))\n",
    "print(f'Files in {ADAPTER_REPO_ID}:')\n",
    "for f in files:\n",
    "    print(f'  {f}')\n",
    "\n",
    "bam_ok = any('bambara/adapter_config.json' in f for f in files)\n",
    "ful_ok = any('fula/adapter_config.json' in f for f in files)\n",
    "print()\n",
    "print(f'Bambara adapter: {\"✅\" if bam_ok else \"❌\"}')\n",
    "print(f'Fula    adapter: {\"✅\" if ful_ok else \"❌\"}')\n",
    "\n",
    "if bam_ok and ful_ok:\n",
    "    print('\\n🎉 Both adapters ready. Your Space will use them automatically on the next reload.')\n",
    "else:\n",
    "    print('\\n⚠️  Some adapters are missing — check the training cells above for errors.')"
   ]
  }
 ],
 "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
}