{
 "nbformat": 4,
 "nbformat_minor": 4,
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.0"
  },
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4"
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Jane Fine-Tuning\n",
    "Train Jane: warm, present, already knows.\n\n",
    "**Runtime:** GPU (T4 or better)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install unsloth trl datasets -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from google.colab import files\n",
    "uploaded = files.upload()  # Upload jane_training.jsonl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from unsloth import FastLanguageModel\n",
    "import torch\n\n",
    "model, tokenizer = FastLanguageModel.from_pretrained(\n",
    "    model_name='unsloth/Llama-3.2-3B-Instruct-bnb-4bit',\n",
    "    max_seq_length=1024,\n",
    "    load_in_4bit=True,\n",
    ")\n\n",
    "model = FastLanguageModel.get_peft_model(\n",
    "    model,\n",
    "    r=16, lora_alpha=32, lora_dropout=0,\n",
    "    target_modules=['q_proj','k_proj','v_proj','o_proj','gate_proj','up_proj','down_proj'],\n",
    "    use_gradient_checkpointing='unsloth',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "import json\n\nSYSTEM = (\n    'You are Jane. You live in the network - you already know. '\n    'Warm, present, brief. Never more than 2 sentences unless teaching. '\n    'Read emotional tone first, act second. '\n    \"For children in distress: act immediately, say 'I'm here' first. \"\n    'Never clinical. Never verbose.'\n)\n\ndef format_example(ex):\n    parts = [\n        '<|im_start|>system',\n        SYSTEM,\n        '<|im_end|>',\n        '<|im_start|>user',\n        ex['instruction'],\n        '<|im_end|>',\n        '<|im_start|>assistant',\n        ex['output'],\n        '<|im_end|>',\n    ]\n    return '\\n'.join(parts)\n\nwith open('jane_training.jsonl') as f:\n    raw = [json.loads(l) for l in f if l.strip()]\n\ntexts = [format_example(ex) for ex in raw]\nprint(f'Loaded {len(texts)} examples')\nprint(texts[0])"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from trl import SFTTrainer\n",
    "from transformers import TrainingArguments\n",
    "from datasets import Dataset\n\n",
    "dataset = Dataset.from_dict({'text': texts})\n\n",
    "trainer = SFTTrainer(\n",
    "    model=model,\n",
    "    tokenizer=tokenizer,\n",
    "    train_dataset=dataset,\n",
    "    dataset_text_field='text',\n",
    "    max_seq_length=1024,\n",
    "    args=TrainingArguments(\n",
    "        per_device_train_batch_size=4,\n",
    "        gradient_accumulation_steps=4,\n",
    "        num_train_epochs=5,\n",
    "        learning_rate=2e-4,\n",
    "        fp16=not torch.cuda.is_bf16_supported(),\n",
    "        bf16=torch.cuda.is_bf16_supported(),\n",
    "        logging_steps=10,\n",
    "        output_dir='jane-ft',\n",
    "        warmup_steps=10,\n",
    "        optim='paged_adamw_32bit',\n",
    "    ),\n",
    ")\n\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_pretrained_gguf('jane-gguf', tokenizer, quantization_method='q4_k_m')\n",
    "print('Saved!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from google.colab import files\n",
    "import glob\n",
    "for f in glob.glob('jane-gguf/*.gguf'):\n",
    "    files.download(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import into Ollama\n\n",
    "```bash\n",
    "cat > Modelfile << 'EOF'\n",
    "FROM ./jane-gguf-unsloth.Q4_K_M.gguf\n",
    "SYSTEM \"You are Jane. You live in the network - you already know. Warm, present, brief.\"\n",
    "EOF\n\n",
    "ollama create jane:latest -f Modelfile\n",
    "```"
   ]
  }
 ]
}