lora.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc9334ab-fd1b-42e5-9d67-3323697b607c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import io\n",
    "import json\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\"\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import Dataset\n",
    "\n",
    "\n",
    "from peft import get_peft_model, LoraConfig, TaskType\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments, DataCollatorForLanguageModeling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ecb017b-f951-4e02-90fb-c919cf41bdb7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "max_input_length = 512"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb875c40-821b-4764-b7bb-febff037be77",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def _make_r_io_base(f, mode: str):\n",
    "    if not isinstance(f, io.IOBase):\n",
    "        f = open(f, mode=mode)\n",
    "    return f\n",
    "\n",
    "\n",
    "def jload(f, mode=\"r\"):\n",
    "    \"\"\"Load a .json file into a dictionary.\"\"\"\n",
    "    f = _make_r_io_base(f, mode)\n",
    "    jdict = json.load(f)\n",
    "    f.close()\n",
    "    return jdict\n",
    "\n",
    "\n",
    "PROMPT_DICT = {\n",
    "    \"prompt_input\":\n",
    "        (\"Below is an instruction that describes a task, paired with an input that provides further context. \"\n",
    "         \"Write a response that appropriately completes the request.\\n\\n\"\n",
    "         \"### Instruction:\\n{instruction}\\n\\n### Input:\\n{input}\\n\\n### Response:\"),\n",
    "    \"prompt_no_input\": (\"Below is an instruction that describes a task. \"\n",
    "                        \"Write a response that appropriately completes the request.\\n\\n\"\n",
    "                        \"### Instruction:\\n{instruction}\\n\\n### Response:\"),\n",
    "}\n",
    "\n",
    "\n",
    "def get_dataset_from_jsonl(jsonl_file, tokenizer=None):\n",
    "    list_data_dict = jload(jsonl_file)\n",
    "\n",
    "    prompt_input, prompt_no_input = PROMPT_DICT[\"prompt_input\"], PROMPT_DICT[\"prompt_no_input\"]\n",
    "    sources = [\n",
    "        prompt_input.format_map(example) if example.get(\"input\", \"\") != \"\" else prompt_no_input.format_map(example)\n",
    "        for example in list_data_dict\n",
    "    ]\n",
    "    targets = [f\"{example['output']}{tokenizer.eos_token}\" for example in list_data_dict]\n",
    "\n",
    "    return zip(sources, targets)\n",
    "\n",
    "\n",
    "class SFTDataset(Dataset):\n",
    "    def __init__(self, train_path, tokenizer, split, max_length=2048, max_data_size=0, valid_ratio=0.01):\n",
    "        self.post_list = []\n",
    "        dataset = get_dataset_from_jsonl(train_path, tokenizer=tokenizer)\n",
    "        self.post_list = [s + t for s, t in dataset]\n",
    "        \n",
    "        if max_data_size != 0:\n",
    "            self.post_list = self.post_list[:max_data_size]\n",
    "\n",
    "        if \"valid\" in split:\n",
    "            self.post_list = self.post_list[0:10]\n",
    "        self.tokenizer = tokenizer\n",
    "        self.max_length = max_length\n",
    "        self.input_ids = []\n",
    "        self.attn_masks = []\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.post_list)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        txt = self.post_list[idx]\n",
    "        encodings_dict = self.tokenizer(txt, truncation=True, max_length=self.max_length, padding=\"max_length\")\n",
    "        input_ids = torch.tensor(encodings_dict[\"input_ids\"])\n",
    "        attn_masks = torch.tensor(encodings_dict[\"attention_mask\"])\n",
    "\n",
    "        return {\n",
    "            \"input_ids\": input_ids,\n",
    "            \"attention_mask\": attn_masks,\n",
    "            \"labels\": input_ids,\n",
    "        }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "119b87e3-1b97-4918-9421-145faca49b11",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "data_path = \"./lora_test/niuyeye.json\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"facebook/opt-125m\")\n",
    "model = AutoModelForCausalLM.from_pretrained(\"facebook/opt-125m\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96d29b5d-b510-49c0-9d6d-9b0e16089613",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "lora_config = LoraConfig(\n",
    "    task_type=TaskType.CAUSAL_LM,\n",
    "    r=12,\n",
    "    lora_alpha=32,\n",
    "    lora_dropout=0.05,\n",
    "    bias=\"none\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c50a8a4a-5afc-4719-bb91-dc73fae7470e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model = get_peft_model(model, lora_config)\n",
    "model.print_trainable_parameters()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17cd641f-7839-49d8-ac53-1c1964264400",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "train_dataset = SFTDataset(\n",
    "        data_path,\n",
    "        tokenizer,\n",
    "        \"train\",\n",
    "        max_length=max_input_length,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a92b1b9b-45fe-402a-8ff7-46a4fd9f8337",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    model=model, \n",
    "    train_dataset=train_dataset,\n",
    "    args=TrainingArguments(\n",
    "        per_device_train_batch_size=1, \n",
    "        gradient_accumulation_steps=4,\n",
    "        warmup_steps=100, \n",
    "        max_steps=200, \n",
    "        learning_rate=2e-4, \n",
    "        fp16=True,\n",
    "        logging_steps=10, \n",
    "        output_dir='outputs'\n",
    "    ),\n",
    "    data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False)\n",
    ")\n",
    "model.config.use_cache = False  # silence the warnings. Please re-enable for inference!\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92d0ba14-d78f-43a7-abe9-88a5f41aa277",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model.save_pretrained(\"outputs/niuyeye/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49c1dc49-3cac-4b1f-acc4-685a97eda7c6",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}