kaggle-training-completion.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5"
   },
   "outputs": [],
   "source": [
    "!pip install -U bitsandbytes\n",
    "!pip install -U peft\n",
    "!pip install -U transformers\n",
    "!pip install -U trl\n",
    "!pip install -U scikit-learn\n",
    "!pip install -U sentencepiece\n",
    "!pip install -U protobuf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import torch\n",
    "from sklearn.model_selection import train_test_split\n",
    "from datasets import Dataset\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TrainingArguments\n",
    "from trl import SFTTrainer, DataCollatorForCompletionOnlyLM\n",
    "from peft import LoraConfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "exp_name = \"gemma\"\n",
    "data_path = \"25k_utext_uprompt.csv\"\n",
    "model_path = \"google/gemma-1.1-7b-it\"\n",
    "output_path = f\"outputs\"\n",
    "model_save_path = f\"{exp_name}_adapter\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = 3\n",
    "batch_size = 1 \n",
    "max_seq_length = 1536\n",
    "lr = 2e-4\n",
    "access_token = \"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\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv(data_path)\n",
    "train_df, val_df = train_test_split(df, test_size=0.1, random_state=42)\n",
    "train_df = train_df.reset_index(drop=True)\n",
    "val_df = val_df.reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds = Dataset.from_pandas(train_df)\n",
    "val_ds = Dataset.from_pandas(val_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(model_path, token=access_token)\n",
    "tokenizer.pad_token = tokenizer.eos_token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_quant_type=\"nf4\",\n",
    "    bnb_4bit_compute_dtype=\"float16\",\n",
    "    bnb_4bit_use_double_quant=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_path,\n",
    "    quantization_config=bnb_config,\n",
    "    trust_remote_code=True,\n",
    "    device_map=\"auto\",\n",
    "    token=access_token\n",
    ")\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.config.gradient_checkpointing = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def token_len(text):\n",
    "    tokenized = tokenizer(text, return_length=True)\n",
    "    length = tokenized[\"length\"][0]\n",
    "    return length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"\"\"\n",
    "<bos><start_of_turn>user\n",
    "Given are two texts, the Rewritten Text was rewritten from the Original Text by using large language model and a sentence of prompt. You are trying to understand how the Original Text was transformed into the Rewritten Text.\n",
    "Original Text: {}\n",
    "Rewritten Text: {}\n",
    "You should analyze the changes in style, tone, structure, content, etc. Come up with a prompt that must have been used to guide the transformation from the Original Text to the Rewritten Text. Now return the prompt ONLY in one sentence.<end_of_turn>\n",
    "<start_of_turn>model\n",
    "Prompt: {}\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def formatting_prompts_func(example):\n",
    "    output_texts = []\n",
    "    for i in range(len(example[\"rewritten_text\"])):\n",
    "        ori_text = example[\"original_text\"][i]\n",
    "        rew_text = example[\"rewritten_text\"][i]\n",
    "        rew_prompt = example[\"rewrite_prompt\"][i]\n",
    "        text = prompt.format(ori_text, rew_text, rew_prompt)\n",
    "        if token_len(text) > max_seq_length:\n",
    "            continue\n",
    "        output_texts.append(text)\n",
    "    return output_texts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response_template = \"Prompt:\"\n",
    "collator = DataCollatorForCompletionOnlyLM(response_template=response_template, tokenizer=tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "peft_config = LoraConfig(\n",
    "    r=64,\n",
    "    lora_alpha=16,\n",
    "    lora_dropout=0.05,\n",
    "    bias=\"none\",\n",
    "    task_type=\"CAUSAL_LM\",\n",
    "    target_modules=[\n",
    "        \"q_proj\", \"k_proj\", \"v_proj\",\n",
    "        \"o_proj\", \"gate_proj\"\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "args = TrainingArguments(\n",
    "    output_dir=output_path,\n",
    "    fp16=True,\n",
    "    learning_rate=lr,\n",
    "    optim=\"paged_adamw_8bit\",\n",
    "    num_train_epochs=epochs,\n",
    "    per_device_train_batch_size=batch_size,\n",
    "    per_device_eval_batch_size=batch_size,\n",
    "    gradient_accumulation_steps=8,\n",
    "    evaluation_strategy=\"steps\",\n",
    "    eval_steps=200,\n",
    "    save_strategy=\"steps\",\n",
    "    save_steps=200,\n",
    "    save_total_limit=5,\n",
    "    logging_steps=200,\n",
    "    lr_scheduler_type=\"cosine\",\n",
    "    warmup_ratio=0.1,\n",
    "    weight_decay=0.001,\n",
    "    report_to=\"none\",\n",
    "    load_best_model_at_end=True,\n",
    "    metric_for_best_model=\"eval_loss\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = SFTTrainer(\n",
    "    model=model,\n",
    "    args=args,\n",
    "    max_seq_length=max_seq_length,\n",
    "    train_dataset=train_ds,\n",
    "    eval_dataset=val_ds,\n",
    "    formatting_func=formatting_prompts_func,\n",
    "    data_collator=collator,\n",
    "    peft_config=peft_config,\n",
    "    packing=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.empty_cache()\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.save_model(model_save_path)\n",
    "tokenizer.save_pretrained(model_save_path)"
   ]
  }
 ],
 "metadata": {
  "kaggle": {
   "accelerator": "gpu",
   "dataSources": [
    {
     "databundleVersionId": 7806901,
     "sourceId": 67121,
     "sourceType": "competition"
    },
    {
     "datasetId": 3600418,
     "sourceId": 6572938,
     "sourceType": "datasetVersion"
    },
    {
     "sourceId": 164964691,
     "sourceType": "kernelVersion"
    },
    {
     "modelInstanceId": 8658,
     "sourceId": 10716,
     "sourceType": "modelInstanceVersion"
    },
    {
     "isSourceIdPinned": true,
     "modelInstanceId": 17852,
     "sourceId": 21555,
     "sourceType": "modelInstanceVersion"
    }
   ],
   "dockerImageVersionId": 30674,
   "isGpuEnabled": true,
   "isInternetEnabled": false,
   "language": "python",
   "sourceType": "notebook"
  },
  "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}