kaggle-training-chat-mistral.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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, DataCollatorForLanguageModeling\n",
    "from trl import SFTTrainer\n",
    "from peft import LoraConfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "exp_name = \"mistral\"\n",
    "data_path = \"25k_utext_uprompt.csv\"\n",
    "model_path = \"mistralai/Mistral-7B-Instruct-v0.1\"\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(\n",
    "    model_path, \n",
    "    token=access_token, \n",
    "    add_eos_token=True\n",
    ")\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "tokenizer.padding_side = \"right\""
   ]
  },
  {
   "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": [
    "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",
    "<s>[INST] 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 output the prompt ONLY in one sentence. [/INST]\n",
    "The prompt is: {}</s>\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": [
    "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\", \"up_proj\", \"down_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=15,\n",
    "    logging_steps=200,\n",
    "    lr_scheduler_type=\"linear\",\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",
    "    peft_config=peft_config,\n",
    "    data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=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": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}