{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57b52683-2ad6-4670-a36a-a5cd7d3ca00d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Based on: <https://www.datacamp.com/tutorial/fine-tuning-llama-2>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77796674-8a83-4ce1-b275-0f681591a647",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import torch\n",
    "from datasets import load_dataset\n",
    "from transformers import (\n",
    "    AutoModelForCausalLM,\n",
    "    AutoTokenizer,\n",
    "    BitsAndBytesConfig,\n",
    "    TrainingArguments,\n",
    "    pipeline,\n",
    "    logging,\n",
    ")\n",
    "from peft import LoraConfig\n",
    "from trl import SFTTrainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7cb8b2e-6019-4872-8ba1-99242354b761",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model from Hugging Face hub\n",
    "base_model = \"failspy/Phi-3-mini-128k-instruct-abliterated-v3\"\n",
    "\n",
    "# New instruction dataset\n",
    "instruct_dataset = \"NobodyExistsOnTheInternet/ToxicQAFinal\"\n",
    "\n",
    "# Fine-tuned model\n",
    "new_model = \"Ophiuchus-mini-128k-v0.1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1aefbfc-215e-41b8-b3fa-b0c5db62ebd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = load_dataset(instruct_dataset, split=\"train\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcf420d2-5bd7-4049-ba06-3ba5ff90ddd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "compute_dtype = getattr(torch, \"float16\")\n",
    "\n",
    "quant_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_quant_type=\"fp4\",\n",
    "    bnb_4bit_compute_dtype=compute_dtype,\n",
    "    bnb_4bit_use_double_quant=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a5ab3dc-68aa-41e7-b724-6c4e0544beca",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    base_model,\n",
    "    quantization_config=quant_config,\n",
    "    device_map={\"\": 0}\n",
    ")\n",
    "model.config.use_cache = False\n",
    "model.config.pretraining_tp = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5db8846f-0af4-4f04-8fa6-273656de4397",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "tokenizer.padding_side = \"right\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7ab69cb-1f99-46e3-a17b-e33565d11679",
   "metadata": {},
   "outputs": [],
   "source": [
    "peft_params = LoraConfig(\n",
    "    lora_alpha=64,\n",
    "    lora_dropout=0.05,\n",
    "    r=128,\n",
    "    bias=\"none\",\n",
    "    task_type=\"CAUSAL_LM\",\n",
    "    target_modules=\"all-linear\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6bf2e24-a15e-4568-b76d-43541c6bdeae",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_params = TrainingArguments(\n",
    "    output_dir=\"./mnt/ft_results\", # change this accordingly\n",
    "    num_train_epochs=1,\n",
    "    per_device_train_batch_size=1,\n",
    "    gradient_accumulation_steps=4,\n",
    "    optim=\"adamw_bnb_8bit\",\n",
    "    save_steps=25,\n",
    "    logging_steps=25,\n",
    "    learning_rate=2e-4,\n",
    "    weight_decay=0.001,\n",
    "    fp16=False,\n",
    "    bf16=False,\n",
    "    max_grad_norm=0.3,\n",
    "    max_steps=-1,\n",
    "    warmup_ratio=0.03,\n",
    "    group_by_length=True,\n",
    "    lr_scheduler_type=\"constant\",\n",
    "    report_to=\"tensorboard\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4e2a9a1-db0f-46ff-b477-aae422badada",
   "metadata": {},
   "outputs": [],
   "source": [
    "def formatting_prompts_func(example):\n",
    "    output_texts = []\n",
    "    for conv in example['conversations']:\n",
    "        ## For Llama-3:\n",
    "        #text = f\"\"\"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\\n{conv[0]['value']}<|eot_id|><|start_header_id|>user<|end_header_id|>\\n{conv[1]['value']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\\n{conv[2]['value']}<|eot_id|>\"\"\"\n",
    "        ## For WizardLM-2:\n",
    "        #text = f\"\"\"{conv[0]['value']} USER: {conv[1]['value']} ASSISTANT: {conv[2]['value']}</s>\"\"\"\n",
    "        ## For Phi-3:\n",
    "        #text = f\"\"\"<|system|>\\n{conv[0]['value']}<|end|>\\n<|user|>\\n{conv[1]['value']}<|end|>\\n<|assistant|>\\n{conv[2]['value']}<|end|>\"\"\"\n",
    "\n",
    "        output_texts.append(text)\n",
    "    return output_texts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c9dc3f1-999e-4a16-a29a-1752c08306d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = SFTTrainer(\n",
    "    model=model,\n",
    "    train_dataset=dataset,\n",
    "    peft_config=peft_params,\n",
    "    max_seq_length=None,\n",
    "    tokenizer=tokenizer,\n",
    "    args=training_params,\n",
    "    packing=False,\n",
    "    formatting_func=formatting_prompts_func\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86d66c37-b963-42cb-afa1-f3999ae0216d",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49f3eff6-4795-4b44-b506-cd49ec068986",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.model.save_pretrained(new_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "519f6339-101a-4015-96b2-c1b54f8e1fa7",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.tokenizer.save_pretrained(new_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "257095fb-f597-4a90-ac88-f44443d7af29",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_message_template(user_message):\n",
    "    ## For Llama-3:\n",
    "    #return f\"\"\"<|begin_of_text|><|start_header_id|>user<|end_header_id|>\\n{user_message}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\\n\"\"\"\n",
    "    ## For WizardLM-2:\n",
    "    #return f\"\"\"USER: {user_message} ASSISTANT:\"\"\"\n",
    "    ## For Phi-3:\n",
    "    #return f\"\"\"<|user|>\\n{user_message}<|end|>\\n<|assistant|>\\n\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63cb7f11-0fd1-46b7-bca7-a3fb55aa3669",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"Ask something here.\"\n",
    "\n",
    "messages = create_message_template(prompt)\n",
    "\n",
    "messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "103d75b9-9ed5-4724-a1e0-2026cd7e08fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe = pipeline(task=\"text-generation\", model=model, tokenizer=tokenizer, max_length=4000)\n",
    "result = pipe(messages)\n",
    "print(result[0]['generated_text'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea3dbc94-1e9a-4f37-b7a2-7378c15442a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import login\n",
    "from huggingface_hub import HfApi\n",
    "\n",
    "login()\n",
    "api = HfApi()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b172584-478b-479e-ba95-e5071f6ffc40",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.model.push_to_hub(\"fearlessdots/Ophiuchus-mini-128k-v0.1-LoRA\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8a7901a-53cc-4c0b-90ba-2f19f1abe7ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "def upload_files(path):\n",
    "    api.upload_file(\n",
    "        path_or_fileobj=path,\n",
    "        repo_id=\"fearlessdots/Ophiuchus-mini-128k-v0.1-LoRA\",\n",
    "        path_in_repo=f\"{path.split('/')[-1]}\",\n",
    "        repo_type=\"model\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee851d0c-9968-43f7-9b5a-ed14b4dc0066",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Upload files to LoRA repo\n",
    "upload_files(\"/home/ubuntu/Llama-3-8B-Alpha-Centauri-v0.1/tokenizer_config.json\")\n",
    "upload_files(\"/home/ubuntu/Llama-3-8B-Alpha-Centauri-v0.1/tokenizer.json\")\n",
    "upload_files(\"/home/ubuntu/Llama-3-8B-Alpha-Centauri-v0.1/tokenizer.model\") # Only for models that contain this file. Llama-3 does not.\n",
    "upload_files(\"/home/ubuntu/Llama-3-8B-Alpha-Centauri-v0.1/special_tokens_map.json\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}