doc/tutorial/inferdpt/inferdpt_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "341aeb6e-9e25-4a0e-9664-a32ab11293fa",
   "metadata": {},
   "source": [
    "# Inferdpt Tutorial"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b40afd5-77b9-45c6-a761-81b9a6bddc05",
   "metadata": {},
   "source": [
    "## Introduction of Inferdpt\n",
    "\n",
    "Inferdpt is an advanced algorithm framework designed for efficient and privacy-preserving text generation using large language models (LLMs). The framework addresses privacy concerns related to data leakage and unauthorized information collection in LLMs. Inferdpt implements Differential Privacy mechanisms to protect sensitive information during the inference process with black-box LLMs.\n",
    "\n",
    "Inferdpt comprises two key modules: the \"perturbation module\" and the \"extraction module\". The perturbation module utilizes a differentially private(DP) mechanism to generate a perturbed prompt from the raw document, facilitating privacy-preserving inference with black-box LLMs. The extraction module, inspired by knowledge distillation and retrieval-augmented generation, processes the perturbed text to produce coherent and consistent output. This ensures that the text generation quality of InferDPT is comparable to that of non-private LLMs, maintaining high utility while providing strong privacy guarantees.\n",
    "\n",
    "To further enhance privacy protection, Inferdpt integrates a novel mechanism called RANTEXT. RANTEXT introduces the concept of random adjacency list for token-level perturbation, addressing the vulnerability of existing differentially private mechanisms to embedding inversion attacks.\n",
    "\n",
    "For more details of Inferdpt, please refer to the [original paper](https://arxiv.org/pdf/2310.12214.pdf)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac982b2d-4a71-45a5-a2b1-90259711f36b",
   "metadata": {},
   "source": [
    "## Use InferDPT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "042049c5-80ce-4786-9896-88baddd59f4e",
   "metadata": {},
   "source": [
    "In this section, we will guide you through the process of:\n",
    "- Setting up the inferdpt toolkit with an existing language model.\n",
    "- Creating a model inference tool using the built-in class.\n",
    "- Executing a step-by-step walkthrough of an inference instance: Employing inferdpt to generate rationale responses for question-answering tasks."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1938eef-106d-4cc0-a9b7-6ad8d9d281f5",
   "metadata": {},
   "source": [
    "### Create Inferdpt Kit"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "565aa2ed-5919-4aa0-9499-23b730434c62",
   "metadata": {},
   "source": [
    "In alignment with the original paper, the implementation of differential privacy in inferdpt involves the random substitution of tokens in the original text with semantically similar words. To facilitate this process, it is necessary to precalculate the similarities between a subset of tokens from the vocabulary of the remote large language model. In this tutorial, we will utilize the Mistral-7B model as our remote large language model and the Qwen1.5-0.5B model as the local decoding model. For the sake of computational efficiency, we will select a subset of 11,400 tokens from the Mistral-7B vocabulary to perform the similarity calculations and use the built-in function to finally get the inferdpt-kit.\n",
    "\n",
    "Firstly we load the mistral model to get the embedding set:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f01a229a-52e1-4a97-af06-a2ab122b7083",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load embeddings from mistral model\n",
    "import numpy as np\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "model_path = '/data/cephfs/llm/models/Mistral-7B-Instruct-v0.2/'\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_path)\n",
    "model = AutoModelForCausalLM.from_pretrained(model_path)\n",
    "embeddings = tokenizer.get_vocab() # get embeddings matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "3f7ec40b-1a58-4608-b2c1-3299979e699a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the embedding layer weights\n",
    "dtype = np.float32\n",
    "embedding_weights = model.get_input_embeddings().weight\n",
    "# Convert the embedding layer weights to numpy\n",
    "embedding_weights_np = embedding_weights.detach().numpy().astype(dtype)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07261aee-b676-4a42-9098-2923fa67519c",
   "metadata": {},
   "source": [
    "Then we select english tokens from the vocabulary. Then we can get an embedding matrix and a corresponding token list."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "dbb231f9-f0ca-4add-bb45-f4fb59429abb",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 32000/32000 [00:00<00:00, 663000.04it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "11400\n"
     ]
    }
   ],
   "source": [
    "import tqdm\n",
    "import re\n",
    "\n",
    "def contains_english_chars(string):\n",
    "    pattern = r'[a-zA-Z]'\n",
    "    match = re.search(pattern, string)\n",
    "    return bool(match)\n",
    "\n",
    "def contains_non_english_chars(string):\n",
    "    pattern = r'[^a-zA-Z]'\n",
    "    match = re.search(pattern, string)\n",
    "    return bool(match)\n",
    "\n",
    "def filter_tokens(token2index):\n",
    "    filtered_index2token = {}\n",
    "    for key, idx in tqdm.tqdm(token2index.items()):\n",
    "        if key.startswith('<'):\n",
    "            continue\n",
    "        if not key.startswith('▁'):\n",
    "            continue\n",
    "        val_ = key.replace(\"▁\", \"\")\n",
    "        if val_ == val_.upper():\n",
    "            continue\n",
    "        if contains_non_english_chars(val_):\n",
    "            continue\n",
    "        if 3 < len(val_) < 16 and contains_english_chars(val_):\n",
    "            filtered_index2token[idx] = key\n",
    "\n",
    "    return filtered_index2token\n",
    "\n",
    "filtered_index2token = filter_tokens(embeddings)\n",
    "used_num_tokens = len(filtered_index2token)\n",
    "print(used_num_tokens)\n",
    "for idx, token in filtered_index2token.items():\n",
    "    token_2_embedding[token] = embedding_weights_np[idx].tolist()\n",
    "token_list = list(token_2_embedding.keys())\n",
    "embedding_matrix = np.array(list(token_2_embedding.values()), dtype=dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "5922a177-d752-485d-98ab-9fd6688198f8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "we got the embedding matrix:\n",
      "[[-6.1035156e-04 -4.5471191e-03 -5.2795410e-03 ... -1.3656616e-03\n",
      "   4.2419434e-03 -8.1634521e-04]\n",
      " [ 4.8522949e-03  5.9814453e-03  1.1596680e-03 ... -2.6702881e-03\n",
      "  -1.7471313e-03  9.9182129e-04]\n",
      " [-2.7465820e-03  4.3029785e-03  3.3874512e-03 ... -2.6092529e-03\n",
      "  -1.2397766e-05 -3.4027100e-03]\n",
      " ...\n",
      " [-6.1340332e-03 -5.3405762e-03 -1.0910034e-03 ... -9.3841553e-04\n",
      "  -7.4005127e-04 -7.3852539e-03]\n",
      " [-4.5166016e-03  8.2015991e-04  4.8217773e-03 ... -1.1978149e-03\n",
      "  -1.0528564e-03 -2.1362305e-03]\n",
      " [ 1.2054443e-03  1.9836426e-03 -2.8419495e-04 ... -1.5792847e-03\n",
      "  -2.8381348e-03 -7.1716309e-04]]\n"
     ]
    }
   ],
   "source": [
    "print('we got the embedding matrix:')\n",
    "print(embedding_matrix)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20890d89-998f-4f38-968a-2a6a0648b050",
   "metadata": {},
   "source": [
    "We can easily prepare the pre-computed data we needed for inferdpt by using the built-in function of the InferDPTKit class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "c90c7099-d20d-4009-bb7a-aeb3b46210b2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "11400it [00:37, 300.99it/s]\n",
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4096/4096 [00:03<00:00, 1147.93it/s]\n"
     ]
    }
   ],
   "source": [
    "from fate_llm.algo.inferdpt.utils import InferDPTKit\n",
    "param = InferDPTKit.make_inferdpt_kit_param(embedding_matrix, token_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fe3a722-5cdf-4393-90f3-e5d7b82051cf",
   "metadata": {},
   "source": [
    "Great, the computation is complete! Now, let’s proceed to perturb a sentence using inferdpt with ε (epsilon) set to 3.0. We will also save the perturbed sentence to a designated folder for future reference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "4a6acd81-bc7c-49d3-86f4-ad0b5c329e61",
   "metadata": {},
   "outputs": [],
   "source": [
    "inferdpt_kit = InferDPTKit(*param, tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "id": "0077696c-0c10-4500-8835-6e72a084bc42",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'into the tree to the woods'"
      ]
     },
     "execution_count": 96,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "inferdpt_kit.perturb('From the river to the ocean', epsilon=3.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "id": "8df2f57a-e202-4d4f-a175-21990223dc3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "save_kit_path = 'your path'\n",
    "inferdpt_kit.save_to_path(save_kit_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ced7f1bf-aa49-4806-92e2-712493bb4b10",
   "metadata": {},
   "source": [
    "### Go through Inferdpt Step by Step\n",
    "\n",
    "Next, we will guide you through the process of using inferdpt step by step. We will simulate the interaction between the client and server locally. Before we begin, let’s discuss model inference. Within fate-llm's inferdpt module, we provide three types of model inference classes: vllm, vllm server, and Huggingface native. You can explore these classes in the [code files](../../../python/fate_llm/algo/inferdpt/inference/) or develop your own inference tool based on your specific needs. We highly recommend using vllm server. In this case, we will use the following two commands to launch two large model services, corresponding to the server’s LLM and the local decoding small model.\n",
    "\n",
    "For this example, we have executed the process on a machine equipped with four V100-32G GPUs. We advise you to adjust the model path and GPU settings as necessary to accommodate the specifications of your own machine.\n",
    "\n",
    "Start vllm server using commands below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1b6c3c7-6ddd-4386-8700-c95f74a2bae0",
   "metadata": {},
   "outputs": [],
   "source": [
    "! python -m vllm.entrypoints.openai.api_server --host 127.0.0.1 --port 8888 --model ./Mistral-7B-Instruct-v0.2  --dtype=half --enforce-eager --tensor-parallel-size 4 --gpu-memory-utilization 0.6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48374cb5-3a5b-456c-9d53-219c2468da63",
   "metadata": {},
   "outputs": [],
   "source": [
    "! python -m vllm.entrypoints.openai.api_server --host 127.0.0.1 --port 8887 --model ./Qwen1.5-0.5B  --dtype=half --enforce-eager --tensor-parallel-size 4 --gpu-memory-utilization 0.2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "375e3f0d-36c7-4ab3-8e65-cccac23e93c6",
   "metadata": {},
   "source": [
    "Next, we will initialize the inference instance, which are the parameters for both the inferdpt client and server. This includes specifying the IP address, port, and the model name of the service that has been started."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "id": "cd099ef4-569d-45b6-9765-502b688c3fb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from fate_llm.inference.api import APICompletionInference\n",
    "# for client\n",
    "inference_client = APICompletionInference(api_url=\"http://127.0.0.1:8887/v1\", model_name='./Qwen1.5-0.5B', api_key='EMPTY')\n",
    "# for server\n",
    "inference_server = APICompletionInference(api_url=\"http://127.0.0.1:8888/v1\", model_name='./Mistral-7B-Instruct-v0.2', api_key='EMPTY')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 135,
   "id": "8c430c14-2180-4f02-8f06-3f41bae1a710",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " I am a new user of this forum. I am a 20 year\n"
     ]
    }
   ],
   "source": [
    "ret = inference_client.inference(['Hello how are you?'], inference_kwargs={\n",
    "    'stop': ['<|im_end|>', '\\n'],\n",
    "    'temperature': 0.01,\n",
    "    'max_tokens': 16\n",
    "})\n",
    "print(ret[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 138,
   "id": "6341eb48-e30f-46d4-aeaa-8c6fd27259b9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " I am an artificial intelligence designed to assist with information and answer questions to the best of my ability. I don't have the ability to have a personal identity or emotions. I'm here to help you with any inquiries you may have. How can I assist you today?\n"
     ]
    }
   ],
   "source": [
    "ret = inference_server.inference(['<s>[INST]Who are u?[/INST]'], inference_kwargs={\n",
    "    'stop': ['</s>'],\n",
    "    'temperature': 0.01,\n",
    "    'max_tokens': 128\n",
    "})\n",
    "print(ret[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90f5ce55-de3b-481a-a9cc-cb4c24edb7c2",
   "metadata": {},
   "source": [
    "In this tutorial, we will use a question-answering (QA) task as our illustrative example. To do so, we will extract a sample from the ARC-E dataset for demonstration purposes, here is the example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "id": "f912f986-ae86-4d57-9ebc-534d6404173c",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_example = {'id': 'Mercury_7220990',\n",
    "'question': 'Which factor will most likely cause a person to develop a fever?',\n",
    "'choices': {'text': ['a leg muscle relaxing after exercise',\n",
    "'a bacterial population in the bloodstream',\n",
    "'several viral particles on the skin',\n",
    "'carbohydrates being digested in the stomach'],\n",
    "'label': ['A', 'B', 'C', 'D']},\n",
    "'answerKey': 'B'}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a98c74a8-7760-438b-b4f4-33178fed8761",
   "metadata": {},
   "source": [
    "Before initiating the inference, it's crucial to understand the sequence of steps involved. We will leverage the Jinja2 template engine to structure our documentation as follows:\n",
    "\n",
    "1. **Document Template Organization**: The initial step is to organize the document dictionary using the DOC TEMPLATE. This template will provide the structure for the input document.\n",
    "\n",
    "2. **Differential Privacy Perturbation**: Apply Differential Privacy (DP) to perturb the structured document string. This will result in a perturbed document. The perturbed document is then added to the original document under the key 'perturbed_doc'. Note that you can modify this key according to your parameter settings.\n",
    "\n",
    "3. **Instruction Addition**: Use the INSTRUCTION TEMPLATE to add instructions (or few-shot examples) to the perturbed document. This modified document is then sent to the server side for processing. The server's response is captured, and this perturbed response is appended to the original document under the key 'perturbed_response'. As before, this key can be adjusted as needed.\n",
    "\n",
    "4. **Decode Template Formatting**: Finally, employ the decode template to format the decode prompt. The resulting inference is then added to the original dictionary under the key 'inferdpt_result'. This key, like the others, can be customized to fit your specific parameters.\n",
    "\n",
    "By following these steps, the inferdpt framework enables a structured and privacy-preserving inference process, leading to a final output that incorporates the perturbed data and the model's response.\n",
    "For more details, you can refer to the source codes:\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09d7377a-22f4-4d04-b886-88faa1384d7f",
   "metadata": {},
   "source": [
    "The templates for this example are defined on the client side. Below is the Jinja template we use:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 141,
   "id": "eff74a65-f765-483f-a685-418376414ff0",
   "metadata": {},
   "outputs": [],
   "source": [
    "doc_template = \"\"\"{{question}} \n",
    "Choices:{{choices.text}}\n",
    "\"\"\"\n",
    "\n",
    "instruction_template=\"\"\"\n",
    "<s>[INST]\n",
    "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:George wants to warm his hands quickly by rubbing them. Which skin surface will produce the most heat?\n",
    "Choices:['dry palms', 'wet palms', 'palms covered with oil', 'palms covered with lotion']\n",
    "Rationale:Friction between two surfaces generates heat due to the conversion of kinetic energy into thermal energy. Dry palms produce the most heat when rubbed together as they create higher friction compared to wet or lubricated palms, which reduce friction.  Therefore, the answer is 'dry palms'.<end>\n",
    "\n",
    "Please explain:\n",
    "Question:{{perturbed_doc}}\n",
    "Rationale:\n",
    "[/INST]\n",
    "\"\"\"\n",
    "\n",
    "decode_template = \"\"\"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:George wants to warm his hands quickly by rubbing them. Which skin surface will produce the most heat?\n",
    "Choices:['dry palms', 'wet palms', 'palms covered with oil', 'palms covered with lotion']\n",
    "Rationale:Friction between two surfaces generates heat due to the conversion of kinetic energy into thermal energy. Dry palms produce the most heat when rubbed together as they create higher friction compared to wet or lubricated palms, which reduce friction.  Therefore, the answer is 'dry palms'.<end>\n",
    "\n",
    "Question:{{perturbed_doc}}\n",
    "Rationale:{{perturbed_response | replace('\\n', '')}}<end>\n",
    "\n",
    "Please explain:\n",
    "Question:{{question}} \n",
    "Choices:{{choices.text}}\n",
    "Rationale:\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3c02898-91df-48b6-a0e2-9af5bd5538d8",
   "metadata": {},
   "source": [
    "Please be aware that we have included a one-shot example in the prompt to ensure that the Large Language Model (LLM) responds as anticipated.\n",
    "\n",
    "Now we create two script: \n",
    "- inferdpt_client.py\n",
    "- inferdpt_server.py\n",
    "\n",
    "And run codes provided below:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0f110a3-c601-4c7b-8e89-8684d2ae266d",
   "metadata": {},
   "source": [
    "#### Client Side: inferdpt_client.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "006612b8-da8d-402c-9b6d-b6786325fa7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from fate_llm.inference.api import APICompletionInference\n",
    "from fate_llm.algo.inferdpt import inferdpt\n",
    "from fate_llm.algo.inferdpt.utils import InferDPTKit\n",
    "import sys\n",
    "\n",
    "\n",
    "arbiter = (\"arbiter\", 10000)\n",
    "guest = (\"guest\", 10000)\n",
    "host = (\"host\", 9999)\n",
    "name = \"fed1\"\n",
    "\n",
    "\n",
    "def create_ctx(local):\n",
    "    from fate.arch import Context\n",
    "    from fate.arch.computing.backends.standalone import CSession\n",
    "    from fate.arch.federation.backends.standalone import StandaloneFederation\n",
    "    import logging\n",
    "\n",
    "    logger = logging.getLogger()\n",
    "    logger.setLevel(logging.INFO)\n",
    "\n",
    "    console_handler = logging.StreamHandler()\n",
    "    console_handler.setLevel(logging.INFO)\n",
    "\n",
    "    formatter = logging.Formatter(\"%(asctime)s - %(name)s - %(levelname)s - %(message)s\")\n",
    "    console_handler.setFormatter(formatter)\n",
    "\n",
    "    logger.addHandler(console_handler)\n",
    "    computing = CSession(data_dir=\"./session_dir\")\n",
    "    return Context(computing=computing, federation=StandaloneFederation(computing, name, local, [guest, host, arbiter]))\n",
    "\n",
    "\n",
    "ctx = create_ctx(guest)\n",
    "save_kit_path = 'your path'\n",
    "kit = InferDPTKit.load_from_path(save_kit_path)\n",
    "inference = APICompletionInference(api_url=\"http://127.0.0.1:8887/v1\", model_name='./Qwen1.5-0.5B', api_key='EMPTY')\n",
    "\n",
    "test_example = {'id': 'Mercury_7220990',\n",
    "'question': 'Which factor will most likely cause a person to develop a fever?',\n",
    "'choices': {'text': ['a leg muscle relaxing after exercise',\n",
    "'a bacterial population in the bloodstream',\n",
    "'several viral particles on the skin',\n",
    "'carbohydrates being digested in the stomach'],\n",
    "'label': ['A', 'B', 'C', 'D']},\n",
    "'answerKey': 'B'}\n",
    "\n",
    "\n",
    "doc_template = \"\"\"{{question}} \n",
    "Choices:{{choices.text}}\n",
    "\"\"\"\n",
    "\n",
    "instruction_template=\"\"\"\n",
    "<s>[INST]\n",
    "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:George wants to warm his hands quickly by rubbing them. Which skin surface will produce the most heat?\n",
    "Choices:['dry palms', 'wet palms', 'palms covered with oil', 'palms covered with lotion']\n",
    "Rationale:Friction between two surfaces generates heat due to the conversion of kinetic energy into thermal energy. Dry palms produce the most heat when rubbed together as they create higher friction compared to wet or lubricated palms, which reduce friction.  Therefore, the answer is 'dry palms'.<end>\n",
    "\n",
    "Please explain:\n",
    "Question:{{perturbed_doc}}\n",
    "Rationale:\n",
    "[/INST]\n",
    "\"\"\"\n",
    "\n",
    "decode_template = \"\"\"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:George wants to warm his hands quickly by rubbing them. Which skin surface will produce the most heat?\n",
    "Choices:['dry palms', 'wet palms', 'palms covered with oil', 'palms covered with lotion']\n",
    "Rationale:Friction between two surfaces generates heat due to the conversion of kinetic energy into thermal energy. Dry palms produce the most heat when rubbed together as they create higher friction compared to wet or lubricated palms, which reduce friction.  Therefore, the answer is 'dry palms'.<end>\n",
    "\n",
    "Question:{{perturbed_doc}}\n",
    "Rationale:{{perturbed_response | replace('\\n', '')}}<end>\n",
    "\n",
    "Please explain:\n",
    "Question:{{question}} \n",
    "Choices:{{choices.text}}\n",
    "Rationale:\n",
    "\"\"\"\n",
    "\n",
    "inferdpt_client = inferdpt.InferDPTClient(ctx, kit, inference, epsilon=3.0)\n",
    "result = inferdpt_client.inference([test_example], doc_template, instruction_template, decode_template, \\\n",
    "                                 remote_inference_kwargs={\n",
    "                                    'stop': ['<\\s>'],\n",
    "                                    'temperature': 0.01,\n",
    "                                    'max_tokens': 256\n",
    "                                 },\n",
    "                                 local_inference_kwargs={\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",
    "print('result is {}'.format(result[0]['inferdpt_result']))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6ed3c0e-0b1f-4087-b155-def3ee957618",
   "metadata": {},
   "source": [
    "#### Server Side: inferdpt_server.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96e3e9fa-9554-4bcf-b8bf-358c469014bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from fate_llm.algo.inferdpt.inferdpt import InferDPTServer\n",
    "import sys\n",
    "from fate_llm.algo.inferdpt.inference.api import APICompletionInference\n",
    "\n",
    "\n",
    "arbiter = (\"arbiter\", 10000)\n",
    "guest = (\"guest\", 10000)\n",
    "host = (\"host\", 9999)\n",
    "name = \"fed1\"\n",
    "\n",
    "\n",
    "def create_ctx(local):\n",
    "    from fate.arch import Context\n",
    "    from fate.arch.computing.backends.standalone import CSession\n",
    "    from fate.arch.federation.backends.standalone import StandaloneFederation\n",
    "    import logging\n",
    "\n",
    "    logger = logging.getLogger()\n",
    "    logger.setLevel(logging.INFO)\n",
    "\n",
    "    console_handler = logging.StreamHandler()\n",
    "    console_handler.setLevel(logging.INFO)\n",
    "\n",
    "    formatter = logging.Formatter(\"%(asctime)s - %(name)s - %(levelname)s - %(message)s\")\n",
    "    console_handler.setFormatter(formatter)\n",
    "\n",
    "    logger.addHandler(console_handler)\n",
    "    computing = CSession(data_dir=\"./session_dir\")\n",
    "    return Context(computing=computing, federation=StandaloneFederation(computing, name, local, [guest, host, arbiter]))\n",
    "\n",
    "\n",
    "ctx = create_ctx(arbiter)\n",
    "inference_server = APICompletionInference(api_url=\"http://127.0.0.1:8888/v1\", model_name='./Mistral-7B-Instruct-v0.2', api_key='EMPTY')\n",
    "inferdpt_server = InferDPTServer(ctx, inference)\n",
    "inferdpt_server.inference()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bfef704b-179e-44cd-84dc-a40b036e7f28",
   "metadata": {},
   "source": [
    "Start two terminal and launch client&server scripts simultaneously.\n",
    "On the client side we can get the answer:\n",
    "\n",
    "```\n",
    "The given question asks which factor will most likely cause a person to develop a fever. The factors mentioned are a leg muscle relaxing after exercise, a bacterial population in the bloodstream, several viral particles on the skin, and carbohydrates being digested in the stomach. The question is asking which factor is most likely to cause a person to develop a fever. The factors are all related to the body's internal environment, but the most likely factor is a bacterial population in the bloodstream. This is because bacteria can cause a fever, and the body's immune system responds to the infection by producing antibodies that can fight off the bacteria. Therefore, the answer is 'a bacterial population in the bloodstream'\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adf80b4e-4727-4ee4-b0b1-4839bd516f4f",
   "metadata": {},
   "source": [
    "## Use Inferdpt in FATE Pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9b560e3-8db4-4828-a4fc-494320a9a3e5",
   "metadata": {},
   "source": [
    "We can leverage the FATE pipeline to submit inference tasks for industrial applications. When operating in pipeline mode, to safeguard against privacy breaches such as API key or server path leakage, it is crucial to create initialization scripts for establishing inferdpt client instances. Alternatively, you can modify the provided scripts within the fate_llm/algo/inferdpt/init folder.\n",
    "\n",
    "Below, we provide an overview of the default_init.py script, which serves as an example of how to create an [initialization class](../../../python/fate_llm/algo/inferdpt/init/default_init.py). By customizing the static variables within this class, you can configure the client and server to interact with the Large Language Model (LLM) interfaces as intended."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eab49960-0541-4059-b84d-bee4bb690974",
   "metadata": {},
   "outputs": [],
   "source": [
    "from fate_llm.algo.inferdpt.init._init import InferClientInit\n",
    "from fate_llm.inference.api import APICompletionInference\n",
    "from fate_llm.algo.inferdpt import inferdpt\n",
    "from fate_llm.algo.inferdpt.utils import InferDPTKit\n",
    "from fate_llm.algo.inferdpt.inferdpt import InferDPTClient, InferDPTServer\n",
    "\n",
    "\n",
    "class InferDPTAPIClientInit(InferClientInit):\n",
    "\n",
    "    api_url = ''\n",
    "    api_model_name = ''\n",
    "    api_key = 'EMPTY'\n",
    "    inferdpt_kit_path = ''\n",
    "    eps = 3.0\n",
    "\n",
    "    def __init__(self, ctx):\n",
    "        super().__init__(ctx)\n",
    "        self.ctx = ctx\n",
    "\n",
    "    def get_inst(self)-> InferDPTClient:\n",
    "        inference = APICompletionInference(api_url=self.api_url, model_name=self.api_model_name, api_key=self.api_key)\n",
    "        kit = InferDPTKit.load_from_path(self.inferdpt_kit_path)\n",
    "        inferdpt_client = inferdpt.InferDPTClient(self.ctx, kit, inference, epsilon=self.eps)\n",
    "        return inferdpt_client\n",
    "\n",
    "\n",
    "class InferDPTAPIServerInit(InferClientInit):\n",
    "\n",
    "    api_url = ''\n",
    "    api_model_name = ''\n",
    "    api_key = 'EMPTY'\n",
    "\n",
    "    def __init__(self, ctx):\n",
    "        super().__init__(ctx)\n",
    "        self.ctx = ctx\n",
    "\n",
    "    def get_inst(self)-> InferDPTServer:\n",
    "        inference = APICompletionInference(api_url=self.api_url, model_name=self.api_model_name, api_key=self.api_key)\n",
    "        inferdpt_server = inferdpt.InferDPTServer(self.ctx,inference_inst=inference)\n",
    "        return inferdpt_server\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a5c9d6b-94b9-4ae3-80f7-20d1a698764c",
   "metadata": {},
   "source": [
    "In the pipeline example, we use arc_easy dataset and our built-in huggingface dataset. Only HuggingfaceDataset is supported in the pipeline mode:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "15276057-fdda-4cc6-8678-eb1f485e4c58",
   "metadata": {},
   "outputs": [],
   "source": [
    "from fate_llm.dataset.hf_dataset import HuggingfaceDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "31cce967-3f5f-4261-ae17-9089368b82f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "dataset = load_dataset('arc_easy')\n",
    "dataset.save_to_disk('your_path/arc_easy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "af9adcb4-766d-45f6-a13c-5c127df61e5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = HuggingfaceDataset(load_from_disk= True, data_split_key='train')\n",
    "ds.load('your_path/arc_easy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "899c410f-fe68-4f7e-936e-8b11720ff148",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'id': 'Mercury_7220990', 'question': 'Which factor will most likely cause a person to develop a fever?', 'choices': {'text': ['a leg muscle relaxing after exercise', 'a bacterial population in the bloodstream', 'several viral particles on the skin', 'carbohydrates being digested in the stomach'], 'label': ['A', 'B', 'C', 'D']}, 'answerKey': 'B'}\n"
     ]
    }
   ],
   "source": [
    "print(ds[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f69f8cf-f40d-418a-be2a-753d67537442",
   "metadata": {},
   "source": [
    "After that, we can associate the dataset path with a name and namespace. By specifying the dataset configuration, the HuggingfaceDataset will be initialized and the dataset will be loaded from the specified path. \n",
    "```\n",
    "flow table bind --namespace experiment --name arc_e --path 'your_path/arc_easy'\n",
    "```\n",
    "Once these initialization scripts are in place, you can submit a pipeline task by specifying the initialization class in the configuration files. For more information, refer to the script provided below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4da1aea7-0ba2-4ebb-918f-cfcf24d4498b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "from fate_client.pipeline.utils import test_utils\n",
    "from fate_client.pipeline.components.fate.reader import Reader\n",
    "from fate_client.pipeline import FateFlowPipeline\n",
    "\n",
    "\n",
    "def main(config=\"../../config.yaml\", namespace=\"\"):\n",
    "    # obtain config\n",
    "    if isinstance(config, str):\n",
    "        config = test_utils.load_job_config(config)\n",
    "    parties = config.parties\n",
    "    guest = parties.guest[0]\n",
    "    arbiter = parties.arbiter[0]\n",
    "\n",
    "    pipeline = FateFlowPipeline().set_parties(guest=guest, arbiter=arbiter)\n",
    "\n",
    "    reader_0 = Reader(\"reader_0\", runtime_parties=dict(guest=guest))\n",
    "    reader_0.guest.task_parameters(\n",
    "        namespace=f\"experiment{namespace}\",\n",
    "        name=\"arc_e\"\n",
    "    )\n",
    "\n",
    "    inferdpt_init_conf_client = {\n",
    "        'module_name': 'fate_llm.algo.inferdpt.init.default_init',\n",
    "        'item_name': 'InferDPTAPIClientInit'\n",
    "    }\n",
    "\n",
    "    dataset_conf = {\n",
    "        'module_name': 'fate_llm.dataset.hf_dataset',\n",
    "        'item_name': 'HuggingfaceDataset',\n",
    "        'kwargs':{\n",
    "            'load_from_disk': True,\n",
    "            'data_split_key': 'train'\n",
    "        }\n",
    "    }\n",
    "\n",
    "    doc_template = \"\"\"{{question}} \n",
    "    Choices:{{choices.text}}\n",
    "    \"\"\"\n",
    "\n",
    "    instruction_template=\"\"\"\n",
    "    <|im_start|>system\n",
    "    You are a helpful assistant.<|im_end|>\n",
    "    <|im_start|>user\n",
    "    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:George wants to warm his hands quickly by rubbing them. Which skin surface will produce the most heat?\n",
    "    Choices:['dry palms', 'wet palms', 'palms covered with oil', 'palms covered with lotion']\n",
    "    Rationale:Friction between two surfaces generates heat due to the conversion of kinetic energy into thermal energy. Dry palms produce the most heat when rubbed together as they create higher friction compared to wet or lubricated palms, which reduce friction.  Therefore, the answer is 'dry palms'.<end>\n",
    "\n",
    "    Please explain:\n",
    "    Question:{{perturbed_doc}}\n",
    "    Rationale:\n",
    "    <|im_end|>\n",
    "    <|im_start|>assistant\n",
    "    \"\"\"\n",
    "\n",
    "    decode_template = \"\"\"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:George wants to warm his hands quickly by rubbing them. Which skin surface will produce the most heat?\n",
    "    Choices:['dry palms', 'wet palms', 'palms covered with oil', 'palms covered with lotion']\n",
    "    Rationale:Friction between two surfaces generates heat due to the conversion of kinetic energy into thermal energy. Dry palms produce the most heat when rubbed together as they create higher friction compared to wet or lubricated palms, which reduce friction.  Therefore, the answer is 'dry palms'.<end>\n",
    "\n",
    "    Question:{{perturbed_doc}}\n",
    "    Rationale:{{perturbed_response | replace('\\n', '')}}<end>\n",
    "\n",
    "    Please explain:\n",
    "    Question:{{question}} \n",
    "    Choices:{{choices.text}}\n",
    "    Rationale:\n",
    "    \"\"\"\n",
    "\n",
    "    remote_inference_kwargs={\n",
    "        'stop': [['<\\s>']],\n",
    "        'temperature': 0.01,\n",
    "        'max_tokens': 256\n",
    "    }\n",
    "\n",
    "    local_inference_kwargs={\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",
    "    inferdpt_client_conf = {\n",
    "        'inferdpt_init_conf': inferdpt_init_conf_client,\n",
    "        'dataset_conf': dataset_conf,\n",
    "        'doc_template': doc_template,\n",
    "        'instruction_template': instruction_template,\n",
    "        'decode_template': decode_template,\n",
    "        'dataset_conf': dataset_conf,\n",
    "        'remote_inference_kwargs': remote_inference_kwargs,\n",
    "        'local_inference_kwargs': local_inference_kwargs\n",
    "    }\n",
    "\n",
    "    inferdpt_init_conf_server = {\n",
    "        'module_name': 'fate_llm.algo.inferdpt.init.default_init',\n",
    "        'item_name': 'InferDPTAPIServerInit'\n",
    "    }\n",
    "\n",
    "    inferdpt_server_conf = {\n",
    "        'inferdpt_init_conf': inferdpt_init_conf_server\n",
    "    }\n",
    "\n",
    "    homo_nn_0 = HomoNN(\n",
    "        'nn_0',\n",
    "        runner_module='inferdpt_runner',\n",
    "        runner_class='InferDPTRunner',\n",
    "        train_data=reader_0.outputs[\"output_data\"]\n",
    "    )\n",
    "\n",
    "    homo_nn_0.guest.task_parameters(runner_conf=inferdpt_client_conf)\n",
    "    homo_nn_0.arbiter.task_parameters(runner_conf=inferdpt_server_conf)\n",
    "    pipeline.add_tasks([reader_0, homo_nn_0])\n",
    "    pipeline.compile()\n",
    "    pipeline.fit()\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    parser = argparse.ArgumentParser(\"PIPELINE DEMO\")\n",
    "    parser.add_argument(\"--config\", type=str, default=\"../config.yaml\",\n",
    "                        help=\"config file\")\n",
    "    parser.add_argument(\"--namespace\", type=str, default=\"\",\n",
    "                        help=\"namespace for data stored in FATE\")\n",
    "    args = parser.parse_args()\n",
    "    main(config=args.config, namespace=args.namespace)\n"
   ]
  }
 ],
 "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.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}