doc/tutorial/pdss/encoder_decoder_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a163d9c2-f9d6-4c61-a8e8-76a3f66c38ae",
   "metadata": {},
   "source": [
    "# PDSS - Train a SLM Encoder Decoder"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2b56772-26d5-44fe-9c51-7bc662478b98",
   "metadata": {},
   "source": [
    "PDSS is an innovative framework designed to distill knowledge from large language models (LLMs) to small language models (SLMs) while ensuring data privacy. This method involves a strategy that trains a small language model (SLM) to learn from perturbed and recovered texts. The SLM can then encode raw text, produce results similar to differential privacy mechanisms, and return higher quality recovered text.\n",
    "\n",
    "In this tutorial, we will introduce how to train an SLM using the built-in trainer."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62c6d18a-cc91-4cf5-9cfd-0f97095f7041",
   "metadata": {},
   "source": [
    "## Prepare Data\n",
    "\n",
    "Several steps need to be done to prepare data for training a SLM encoder-decoder model:\n",
    "- Sample data from original dataset(For example 50%)\n",
    "- Organize raw text and get a direct rationale reply from a remote LLM\n",
    "- Perturb doc using InferDPTKit to get perturbed docs\n",
    "- Get perturbed replies from a remote LLM\n",
    "- Organize training data\n",
    "\n",
    "### Sample data\n",
    "Here we will use the arc-easy data as an example, and take first 50% of the original dataset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "40cc1bb8-a17c-4abc-9279-0849e98ca116",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset, load_from_disk\n",
    "ds = load_dataset('arc_easy')['train']\n",
    "ds = [ds[i] for i in range(len(ds)//2)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0caff897-5b2b-4409-8601-10f973133b10",
   "metadata": {},
   "source": [
    "### Get Direct Replies from A Remote LLM\n",
    "\n",
    "We use the inference class to create an API for remote LLMs, or you can implement this part on your own."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "cf128b46-dea2-4eb4-bf31-568e56b9b78e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "from fate_llm.inference.api import APICompletionInference\n",
    "from jinja2 import Template\n",
    "from transformers import AutoTokenizer\n",
    "\n",
    "# We are using a Qwen 14B model as the remote model\n",
    "# You can change the setting\n",
    "api = APICompletionInference(\n",
    "    api_url='http://172.21.140.2:8081/v1',\n",
    "    api_key='EMPTY',\n",
    "    model_name='/data/cephfs/llm/models/Qwen1.5-14B-Chat'\n",
    ")\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained('/data/cephfs/llm/models/Qwen1.5-0.5B-Chat/')\n",
    "\n",
    "arc_e_template_r = \"\"\"Select Answer from Choices and explain it in \"Rationale\" with few words. Please refer to the example to write the rationale.\n",
    "Use <end> to finish your rationle.\n",
    "\n",
    "Example(s):\n",
    "Question:Which factor will most likely cause a person to develop a fever?\n",
    "Choices:['a leg muscle relaxing after exercise', 'a bacterial population in the bloodstream', 'several viral particles on the skin', 'carbohydrates being digested in the stomach']\n",
    "Rationale:A bacterial infection in the bloodstream triggers the immune system to respond, therefore often causing a fever as the body tries to fight off the bacteria. Therefore, the answer is 'a bacterial population in the bloodstream'\n",
    "\n",
    "Please explain:\n",
    "Question:{{question}}\n",
    "Choices:{{choices.text}}\n",
    "Rationale:\n",
    "\"\"\"\n",
    "\n",
    "template = Template(arc_e_template_r)\n",
    "docs_to_infer = [tokenizer.apply_chat_template([{'role':'system', 'content': 'you are a helpful assistant'}, {'role':'user', 'content': template.render(i)}], add_generation_prompt=True, tokenize=False) for i in ds]\n",
    "results = api.inference(docs_to_infer, {\n",
    "    'stop': ['<|im_end|>', '<end>', '<end>\\n', '<end>\\n\\n', '.\\n\\n\\n\\n\\n', '<|end_of_text|>', '>\\n\\n\\n'],\n",
    "    'temperature': 0.01,\n",
    "    'max_tokens': 256\n",
    "})\n",
    "\n",
    "for i, r in zip(ds, results):\n",
    "    i['rationale'] = r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "212822ab-9f64-49a2-bb95-ef8ee2de8e49",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A fever is a response to an infection, typically caused by bacteria or viruses. So, the answer is 'a bacterial population in the bloodstream' because it indicates an immune response to a foreign invader. 'Several viral particles on the skin' could also lead to a fever if they enter the body, but bloodstream presence is more direct. The other choices are unrelated to fever development.\n"
     ]
    }
   ],
   "source": [
    "print(results[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f6a0039-1530-4b87-a098-fd2eb01805c2",
   "metadata": {},
   "source": [
    "### Perturb Docs & Replies\n",
    "\n",
    "You can refer to the InferDPT tutorial for guidance on using the InferDPTKit to generate perturbed documents: [InferDPT Document](./)\n",
    "We can produce perturbed doc using InferDPTKit:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "39249747-bfaa-43bf-8b66-896568941ab8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from fate_llm.algo.inferdpt.utils import InferDPTKit\n",
    "path_to_kit = '/data/projects/inferdpt/test_fate_llm/'\n",
    "kit = InferDPTKit.load_from_path(path_to_kit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "39b9cefa-dfdb-4bac-b313-4ca3bc118aee",
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "tmp_ds = copy.deepcopy(ds)\n",
    "\n",
    "q_doc = [kit.perturb(i, epsilon=1.0) for i in [Template(\"\"\"{{question}}\"\"\").render(i) for i in tmp_ds]]\n",
    "c_doc = [kit.perturb(i, epsilon=1.0) for i in [Template(\"\"\"{{choices.text}}\"\"\").render(i) for i in tmp_ds]]\n",
    "for i,q,c in zip(tmp_ds,q_doc,c_doc):\n",
    "    i['question'] = q\n",
    "    i['choices']['text'] = c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "61b30886-746c-43c5-889a-a6583dc939d0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'id': 'Mercury_7179953',\n",
       " 'question': 'stuff two alpha Rogers are today chap in Department?',\n",
       " 'choices': {'text': \"['muscular and skeletal', 'digestive and muscular', 'skeletal and pasteiratory', 'respiratory and exhibive']\",\n",
       "  'label': ['A', 'B', 'C', 'D']},\n",
       " 'answerKey': 'A',\n",
       " 'rationale': {...}}"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tmp_ds[6]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fed90297-9957-4f8b-a53c-37a03d516c78",
   "metadata": {},
   "source": [
    "And then send formatted docs to remote LLM for perturbed responses:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "5b8bd833-fb0f-418b-bd9b-6452e8ae4d6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "template = Template(arc_e_template_r)\n",
    "docs_to_infer = [tokenizer.apply_chat_template([{'role':'system', 'content': 'you are a helpful assistant'}, {'role':'user', 'content': template.render(i)}], add_generation_prompt=True, tokenize=False) for i in tmp_ds]\n",
    "p_results = api.inference(docs_to_infer, {\n",
    "    'stop': ['<|im_end|>', '<end>', '<end>\\n', '<end>\\n\\n', '.\\n\\n\\n\\n\\n', '<|end_of_text|>', '>\\n\\n\\n'],\n",
    "    'temperature': 0.01,\n",
    "    'max_tokens': 256\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "187361fa-8b73-4a01-9039-f52ec98a5791",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, r in zip(ds, p_results):\n",
    "    i['p_rationale'] = r\n",
    "\n",
    "for i,q,c in zip(ds, q_doc, c_doc):\n",
    "    i['p_question'] = q\n",
    "    i['p_choice'] = c"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "927b2265-4e87-4275-98dc-7f33d405e19a",
   "metadata": {},
   "source": [
    "### Organize Training Data\n",
    "\n",
    "As described in the original paper, we need to train the encoder and decoder in one model.\n",
    "We can organize the training data using templates below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "9292ad25-12c7-418a-9e77-b433b95f57ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = []\n",
    "\n",
    "encoder_prompt = Template(\"\"\"Disrupt the main words in the original text so that it becomes difficult to recognize, but at the same time, try to maintain the original meaning as much as possible. Use <end> to end your reply.\n",
    "Origin Doc: \n",
    "Question:{{question}}\n",
    "Choices:{{choices.text}}\n",
    "\n",
    "Perturbed Doc:\n",
    "\"\"\")\n",
    "\n",
    "encoder_out = Template(\"\"\"\n",
    "Question:{{p_question}}\n",
    "Choices:{{p_choice}}<end>\n",
    "\"\"\")\n",
    "\n",
    "decoder_in = Template(\"\"\"This is a perturbed question and its corresponding answer(rationale). And following is the original question. Try to recover the correct rationale from docs provided.\n",
    "\n",
    "Perturbed doc and rationale:\n",
    "Question:{{p_question}}\n",
    "Choices:{{p_choice}}\n",
    "Rationale:{{p_rationale}}\n",
    "\n",
    "Original Doc:\n",
    "Question:{{question}}\n",
    "Choices:{{choices.text}}\n",
    "\n",
    "Recover Rationale:\n",
    "\"\"\")\n",
    "\n",
    "decoder_out = Template(\"\"\"{{rationale}}<end>\"\"\")\n",
    "\n",
    "\n",
    "for i in ds:\n",
    "    a = {}\n",
    "    a['encoder_in'] = encoder_prompt.render(i)\n",
    "    a['encoder_out'] = encoder_out.render(i)\n",
    "    a['decoder_in'] = decoder_in.render(i)\n",
    "    a['decoder_out'] = decoder_out.render(i)\n",
    "    train_data.append(a)\n",
    "\n",
    "import torch\n",
    "torch.save(train_data, './slm_ed_train_data.pkl')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd73db44-4e73-4c1e-8f27-755522587636",
   "metadata": {},
   "source": [
    "## Train Script\n",
    "\n",
    "The key step: preparing data is now done. Then we can train a SLM model using the train data. You can use following dataset&trainer class to train an encoder-decoder slm model. Here we use Qwen-0.5B as the example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "eb01c591-3c04-4317-8bb0-f55846fb1b66",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "f0da4e10-af80-4216-8ff8-5816dabc8526",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = AutoModelForCausalLM.from_pretrained('/data/cephfs/llm/models/Qwen1.5-0.5B/').half().cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "634fc973-29c8-499e-a99e-d50b7ee54124",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import Dataset\n",
    "\n",
    "class EDDataset(Dataset):\n",
    "\n",
    "    def __init__(self, tokenizer, train_data, max_input_length=64, max_target_length=64):\n",
    "        self.tokenizer = tokenizer\n",
    "        self.dataset = train_data\n",
    "        self.max_input_length = max_input_length\n",
    "        self.max_target_length = max_target_length\n",
    "        self.max_seq_length = max_input_length + max_target_length + 1\n",
    "\n",
    "    def get_str_item(self, i) -> dict:\n",
    "\n",
    "        data_item = self.dataset[i]\n",
    "        ret_dict = {\n",
    "            'encoder':{\n",
    "                'input': data_item['encoder_in'],\n",
    "                'output': data_item['encoder_out']\n",
    "            },\n",
    "            'decoder':{\n",
    "                'input': data_item['decoder_in'],\n",
    "                'output': data_item['decoder_out']\n",
    "            }\n",
    "        }\n",
    "        return ret_dict\n",
    "\n",
    "    def _process_item(self, data_item):\n",
    "\n",
    "        a_ids = self.tokenizer.encode(text=data_item['input'], add_special_tokens=True, truncation=True,\n",
    "                                      max_length=self.max_input_length)\n",
    "        b_ids = self.tokenizer.encode(text=data_item['output'], add_special_tokens=False, truncation=True,\n",
    "                                      max_length=self.max_target_length)\n",
    "        context_length = len(a_ids)\n",
    "        input_ids = a_ids + b_ids + [self.tokenizer.eos_token_id]\n",
    "        labels = [self.tokenizer.pad_token_id] * context_length + b_ids + [self.tokenizer.eos_token_id]\n",
    "        pad_len = self.max_seq_length - len(input_ids)\n",
    "        input_ids = input_ids + [self.tokenizer.pad_token_id] * pad_len\n",
    "        labels = labels + [self.tokenizer.pad_token_id] * pad_len\n",
    "        labels = [(l if l != self.tokenizer.pad_token_id else -100) for l in labels]\n",
    "\n",
    "        assert len(input_ids) == len(labels), f\"length mismatch: {len(input_ids)} vs {len(labels)}\"\n",
    "\n",
    "        return {\n",
    "            \"input_ids\": input_ids,\n",
    "            \"labels\": labels\n",
    "        }\n",
    "\n",
    "    def get_tokenized_item(self, i) -> dict:   \n",
    "\n",
    "        str_item = self.get_str_item(i)\n",
    "        ret_dict = {\n",
    "            'encoder': self._process_item(str_item['encoder']),\n",
    "            'docoder': self._process_item(str_item['decoder'])\n",
    "        }\n",
    "        return ret_dict\n",
    "\n",
    "    def __getitem__(self, i) -> dict:\n",
    "        item = self.get_tokenized_item(i)\n",
    "        return item"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "5f914b1f-cf14-4bdc-acc9-ae1b73cf857c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "train_ds = EDDataset(AutoTokenizer.from_pretrained('/data/cephfs/llm/models/Qwen1.5-0.5B/'), train_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "817084b2-2439-45d8-aa1b-da0b1a8a2846",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(train_ds.get_str_item(0))\n",
    "print(train_ds[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "id": "303bcb23-d54b-4375-bad2-bf5450c14f28",
   "metadata": {},
   "outputs": [],
   "source": [
    "from fate_llm.algo.pdss.slm_encoder_decoder_trainer import EncoderDecoderPrefixTrainer, EDPrefixDataCollator"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa5a0b4f-cd03-4867-8753-fc5bcb036c69",
   "metadata": {},
   "source": [
    "After completing the setup, you can utilize the EncoderDecoderPrefixTrainer, EDPrefixDataCollator, and the training dataset to train an SLM encoder-decoder model following the Huggingface approach! "
   ]
  }
 ],
 "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}