{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "714f4501",
   "metadata": {},
   "source": [
    "```\n",
    "(notebook) root@a70d062f4f65:~/desktop# cd /\n",
    "(notebook) root@a70d062f4f65:/# jupyter notebook --ip 0.0.0.0 --port 8080 --allow-root\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "cdab2518",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install --upgrade accelerate datasets peft torch transformers trl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5a82e319",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b74c5471",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_id = \"google/gemma-3-1b-it\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8ba9ad8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(model_id, token=os.environ['SIA_HF_API_KEY'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e183a2f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_id,\n",
    "    torch_dtype=torch.bfloat16,\n",
    "    device_map={\"\": 0},\n",
    "    token=os.environ['SIA_HF_API_KEY'],\n",
    "    attn_implementation='eager',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "143138b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "exec(open(\"/root/sia/tools/train/train/dataset.py\").read())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "703eb030",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = Dataset(\"/root/sia/training/config.yaml\")\n",
    "dataset.validate()\n",
    "dataset = dataset.to_transformers_dataset(tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "291f57e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_sia_example(example):\n",
    "    return example['messages'].removeprefix(\"<bos>\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "f01f4125",
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import LoraConfig\n",
    "\n",
    "lora_config = LoraConfig(\n",
    "    r=8,\n",
    "    target_modules=[\"q_proj\", \"o_proj\", \"k_proj\", \"v_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
    "    task_type=\"CAUSAL_LM\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "e2400819",
   "metadata": {},
   "outputs": [],
   "source": [
    "import transformers\n",
    "from trl import SFTTrainer\n",
    "\n",
    "trainer = SFTTrainer(\n",
    "    model=model,\n",
    "    train_dataset=dataset,\n",
    "    args=transformers.TrainingArguments(\n",
    "        per_device_train_batch_size=1,\n",
    "        gradient_accumulation_steps=1,\n",
    "        warmup_steps=1,\n",
    "        max_steps=1,\n",
    "        learning_rate=1e-3,\n",
    "        fp16=True,\n",
    "        logging_steps=1,\n",
    "        output_dir=\"/root/models/notebook_train\",\n",
    "        optim=\"paged_adamw_8bit\"\n",
    "    ),\n",
    "    peft_config=lora_config,\n",
    "    formatting_func=format_sia_example,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "ce61cf71",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "b4e9cc7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = trainer.model.merge_and_unload()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "6c429c9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_pretrained(\"/root/models/notebook/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "ef091e77",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.save_pretrained(\"/root/models/notebook/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "30c823ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "!git clone https://github.com/ggml-org/llama.cpp.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "4a1f4b71",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -r llama.cpp/requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "51edd868",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-cpp-python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a62a2175",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python ./llama.cpp/convert_hf_to_gguf.py --outfile /root/models/notebook.gguf  --outtype q8_0 /root/models/notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fea3615",
   "metadata": {},
   "outputs": [],
   "source": [
    "!apt update && apt install -y unzip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2418ba6",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://github.com/ggml-org/llama.cpp/releases/download/b5226/llama-b5226-bin-ubuntu-x64.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "262270a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "unzip llama-b5226-bin-ubuntu-x64.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff438dcc",
   "metadata": {},
   "outputs": [],
   "source": [
    "!cd build/bin/ && ./llama-cli -m /root/models/notebook.gguf"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "notebook",
   "language": "python",
   "name": "notebook"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}