{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# autoresearch: Coding skill model (SFT, Qwen2.5-Coder-0.5B)\n\nRun this notebook on a Colab GPU runtime (Runtime > Change runtime type > T4). It fine-tunes a small code LLM with QLoRA (4-bit) and pushes the result to **40Hz/autoresearch-coding-v1**. Checkpoints go to Google Drive (survive session timeout) with Trackio logging. See `tracks/train_coding_sft.py` for the same logic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip -q install -U trl peft accelerate datasets bitsandbytes huggingface_hub trackio && echo INSTALLED"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\nfrom google.colab import userdata, drive\ndrive.mount(\"/content/drive\")  # checkpoints survive session timeouts here\nos.environ[\"HF_TOKEN\"] = userdata.get(\"HF_TOKEN\")  # set via Secrets tab\nprint(\"token set:\", bool(os.environ[\"HF_TOKEN\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\nfrom datasets import load_dataset\nfrom peft import LoraConfig\nfrom trl import SFTTrainer, SFTConfig\n\n# evol-codealpaca-v1 (Apache-2.0, GPT-4 evolved, ~111K rows) sliced to 60K for a\n# single T4 epoch; swap to the 18K python set for a lean run.\nDATA_ID = os.environ.get(\"CODING_DATA\", \"theblackcat102/evol-codealpaca-v1\")\nN_ROWS = int(os.environ.get(\"CODING_ROWS\", \"60000\"))\ndataset = load_dataset(DATA_ID, split=\"train\")\nds = dataset.shuffle(seed=42).select(range(N_ROWS))\nds = ds.train_test_split(test_size=0.05, seed=42)\n\ndef fmt(rows):\n    texts = []\n    for i, o in zip(rows[\"instruction\"], rows[\"output\"]):\n        texts.append(f\"### Instruction\\n{i}\\n\\n### Response\\n{o}<|endoftext|>\")\n    return {\"text\": texts}\n\ntrain_ds = ds[\"train\"].map(fmt, batched=True, remove_columns=ds[\"train\"].column_names)\neval_ds = ds[\"test\"].map(fmt, batched=True, remove_columns=ds[\"test\"].column_names)\nprint(\"train/eval:\", len(train_ds), len(eval_ds))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import BitsAndBytesConfig\n",
    "from peft import LoraConfig\n",
    "from trl import SFTTrainer, SFTConfig\n",
    "\n",
    "# QLoRA: 4-bit NF4 with fp16 compute (T4 is Turing; fp16 yes, bf16 no).\n",
    "bnb = BitsAndBytesConfig(\n",
    "    load_in_4bit=True, bnb_4bit_quant_type=\"nf4\",\n",
    "    bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.float16,\n",
    ")\n",
    "# output dir on Drive so checkpoints survive a session timeout; resume with\n",
    "# trainer.train(resume_from_checkpoint=True)\n",
    "OUT = \"/content/drive/MyDrive/autoresearch-coding-v1\"\n",
    "\n",
    "lora = LoraConfig(r=16, lora_alpha=32, target_modules=[\"q_proj\",\"k_proj\",\"v_proj\",\"o_proj\"], lora_dropout=0.05)\n",
    "args = SFTConfig(\n",
    "    output_dir=OUT,\n",
    "    push_to_hub=True,\n",
    "    hub_model_id=\"40Hz/autoresearch-coding-v1\",\n",
    "    num_train_epochs=2,\n",
    "    per_device_train_batch_size=4,\n",
    "    gradient_accumulation_steps=4,\n",
    "    gradient_checkpointing=True,\n",
    "    eval_strategy=\"steps\",\n",
    "    eval_steps=200,\n",
    "    logging_steps=20,\n",
    "    save_steps=200,\n",
    "    save_total_limit=3,\n",
    "    fp16=True,\n",
    "    learning_rate=2e-4,\n",
    "    lr_scheduler_type=\"cosine\",\n",
    "    warmup_ratio=0.05,\n",
    "    optim=\"paged_adamw_8bit\",\n",
    "    max_length=1024,\n",
    "    report_to=\"trackio\",\n",
    "    trackio_space_id=os.environ.get(\"TRACKIO_SPACE_ID\", \"40Hz/autoresearch-gpu\"),\n",
    "    model_init_kwargs={\"quantization_config\": bnb, \"torch_dtype\": torch.float16},\n",
    ")\n",
    "trainer = SFTTrainer(\n",
    "    model=\"Qwen/Qwen2.5-Coder-0.5B\",\n",
    "    train_dataset=train_ds,\n",
    "    eval_dataset=eval_ds,\n",
    "    peft_config=lora,\n",
    "    args=args,\n",
    ")\n",
    "trainer.train()\n",
    "trainer.push_to_hub()\n",
    "print(\"DONE -> 40Hz/autoresearch-coding-v1\")"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}