doc/tutorial/fdkt/fdkt.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Synthesize Data With FDKT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this tutoria, we will demonstrate how to Synthesize data using the FATE-LLM framework. In FATE-LLM, we introduce the \"FDKT\" module,  specifically designed for domain-specific knowledge transfer on large language models using synthetic data. FDKT Algorithm is based on paper [Federated Domain-Specific Knowledge Transfer on\n",
    "Large Language Models Using Synthetic Data](https://arxiv.org/pdf/2405.14212), We integrate its code into the FATE-LLM framework.  \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset: Yelp\n",
    "We processed and sample data of 'Health' subdomain from [Yelp dataset](https://arxiv.org/abs/1509.01626) , the dataset can be downloaded from [here](https://www.yelp.com/dataset). \n",
    "Once the dataset has been downloaded, execute the following command to untar the downloaded dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "```shell\n",
    "tar -xvf yelp_dataset.tar\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code will sample 5000 datalines of 'Health' subdomain, and train data will generated under the folder './processed_data/Health/train.json'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import sys\n",
    "import random\n",
    "from pathlib import Path\n",
    "random.seed(42)\n",
    "\n",
    "\n",
    "base_dir = \"./\"\n",
    "business_data_path = os.path.join(base_dir, 'yelp_academic_dataset_business.json')\n",
    "review_data_path = os.path.join(base_dir, 'yelp_academic_dataset_review.json')\n",
    "\n",
    "business_data_file = open(business_data_path, 'r')\n",
    "review_data_file = open(review_data_path, 'r')\n",
    "\n",
    "categories_list = ['Restaurants', 'Shopping', 'Arts', 'Health']\n",
    "business_dic = {}\n",
    "data_dict = {}\n",
    "for category in categories_list:\n",
    "    business_dic[category] = set()\n",
    "    data_dict[category] = []\n",
    "\n",
    "\n",
    "def get_categories(categories):\n",
    "    return_list = []\n",
    "    for category in categories_list:\n",
    "        if category in categories:\n",
    "            return_list.append(category)\n",
    "    return return_list\n",
    "\n",
    "\n",
    "for line in business_data_file.readlines():\n",
    "    dic = json.loads(line)\n",
    "    if 'categories' in dic.keys() and dic['categories'] is not None:\n",
    "        category = get_categories(dic['categories'])\n",
    "        if len(category) == 1:\n",
    "            business_dic[category[0]].add(dic['business_id'])\n",
    "\n",
    "# for category in categories_list:\n",
    "for line in review_data_file.readlines():\n",
    "    dic = json.loads(line)\n",
    "    if 'business_id' in dic.keys() and dic['business_id'] is not None:\n",
    "        for category in categories_list:\n",
    "            if dic['business_id'] in business_dic[category]:\n",
    "                if dic['text'] is not None and dic['stars'] is not None:\n",
    "                    data_dict[category].append({'text': dic['text'], 'stars': dic['stars']})\n",
    "                break\n",
    "\n",
    "train_data_path = os.path.join('processed_data', \"Health\", 'train.json')\n",
    "os.makedirs(Path(train_data_path).parent, exist_ok=True)\n",
    "train_data_file = open(train_data_path, 'w')\n",
    "data_list = data_dict[\"Health\"]\n",
    "\n",
    "sample_data_dict = dict()\n",
    "\n",
    "for data in data_list:\n",
    "    star = int(data[\"stars\"])\n",
    "    if star not in sample_data_dict:\n",
    "        sample_data_dict[star] = []\n",
    "\n",
    "    sample_data_dict[star].append(data)\n",
    "\n",
    "data_list = []\n",
    "star_keys = list(sample_data_dict.keys())\n",
    "for star in star_keys:\n",
    "    sample_data = sample_data_dict[star][:1000]\n",
    "    random.shuffle(sample_data)\n",
    "    data_list.extend(sample_data)\n",
    "\n",
    "random.shuffle(data_list)\n",
    "json.dump(data_list, train_data_file, indent=4)\n",
    "train_data_file.close()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Models Use\n",
    "Please download the following models, these models are used for data augmentation process.\n",
    "\n",
    "LLM: [Qwen1.5-7B-Chat](https://huggingface.co/Qwen/Qwen1.5-7B-Chat)  \n",
    "SLM: [gpt2-xl](https://huggingface.co/openai-community/gpt2-xl)\n",
    "\n",
    "MeanWhile, 'all-mpnet-base-v2' is used to generate embedding vectors in LLM side.\n",
    "\n",
    "Embedding Model:  [all-mpnet-base-v2](https://huggingface.co/sentence-transformers/all-mpnet-base-v2)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running FDKT Data Synthetic Process With Launcher (Experimential Using)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### SLM Setting\n",
    "\n",
    "In this section, we will introduce some key configurations in SLM side."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1. loading model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import transformers\n",
    "from fate_llm.data.tokenizers.cust_tokenizer import get_tokenizer\n",
    "\n",
    "\n",
    "slm_pretrained_path = \"gpt2-xl\" # modity this to local directory\n",
    "slm = transformers.AutoModelForCausalLM.from_pretrained(slm_pretrained_path, torch_dtype=torch.bfloat16)\n",
    "tokenizer = get_tokenizer(slm_pretrained_path)\n",
    "tokenizer.pad_token_id = tokenizer.eos_token_id\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2. Initialize SLM Training Arugments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fate_llm.algo.fdkt.fdkt_data_aug import FDKTTrainingArguments\n",
    "\n",
    "\n",
    "training_args = FDKTTrainingArguments(\n",
    "    use_cpu=False, # use gpu to do dp(differential privacy) training process\n",
    "    device_id=0, # the device number of gpu\n",
    "    num_train_epochs=1, # dp training epochs\n",
    "    per_device_train_batch_size=2, # batch size of dp training\n",
    "    slm_generation_batch_size=32, # batch_size to generate data in slm side\n",
    "    seq_num_for_single_category=300, # data num for each category(label)\n",
    "    slm_generation_config=dict(\n",
    "        max_new_tokens=256,\n",
    "        temperature=1.0,\n",
    "        top_k=50,\n",
    "        top_p=0.9,\n",
    "        repetition_penalty=1.0,\n",
    "        pad_token_id=tokenizer.eos_token_id\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 3. Initlaize DataSet Instance\n",
    "\n",
    "We provide default templates for dataset \"Yelp\" and \"AGNews\", user can refer [here](https://github.com/FederatedAI/FATE-LLM/tree/dev-2.2.0/python/fate_llm/dataset/data_config) for more details. If you want to use your own dataset, please provide fields label_key/text_key/augment_format/filter_format/tokenize_format/sub_domain/label_list/few_shot_format/text_with_label_format like the two default templates and passing it as and argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fate_llm.dataset.flex_dataset import FlexDataset\n",
    "\n",
    "\n",
    "ds = FlexDataset(\n",
    "    tokenizer_name_or_path=slm_pretrained_path,\n",
    "    load_from=\"json\",\n",
    "    data_part=\"train\",\n",
    "    dataset_name=\"yelp_review\", # use default template\n",
    "    # config=dict/template_path # if dataset_name not equals to \"yelp_review\" or \"ag_news\"\n",
    "    need_preprocess=True,\n",
    "    select_num=2000, # use data_num=2000 to train, default is None, None means using all data\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### LLM Setting\n",
    "\n",
    "In this section, we will introduce some key configurations in LLM side."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1. Deploy VLLM Server And Use OpenAI API Protocol To SpeedUp LLM Inference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "please copy the following code to local file create_and_start_vllm.sh, then run the bash code by executing \"bash create_and_start_vllm.sh\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create_and_start_vllm.sh\n",
    "# create vllm enviroment\n",
    "\n",
    "python -m venv vllm_venv\n",
    "source vllm_venv/bin/activate\n",
    "pip install vllm==0.4.3\n",
    "pip install numpy==1.26.4 # numpy >= 2.0.0 will raise error, so reinstall numpy<2.0.0\n",
    "\n",
    "# please modify Qwen1.5-7B-Chat to local llm model saving path\n",
    "export CUDA_VISIBLE_DEVICES=1,2\n",
    "nohup python -m vllm.entrypoints.openai.api_server --host 127.0.0.1 --port 9999 --model Qwen1.5-7B-Chat --dtype=half --enforce-eager --api-key demo --device cuda -tp 2 &"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2. Initialize LLM Training Arugments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fate_llm.algo.fdkt.fdkt_data_aug import FDKTTrainingArguments\n",
    "\n",
    "\n",
    "training_args = FDKTTrainingArguments(\n",
    "    sample_num_per_cluster=4, # use this to estimate the number of clusters, n_clusters=(len(dataset) + sample_num_per_cluster - 1) // sample_num_per_cluster\n",
    "    filter_prompt_max_length=2**16,\n",
    "    filter_generation_config=dict(\n",
    "        max_tokens=512,\n",
    "    ),\n",
    "    aug_generation_config=dict(\n",
    "        max_tokens=4096,\n",
    "        temperature=0.8,\n",
    "        top_p=0.9,\n",
    "    ),\n",
    "    aug_prompt_num=20000, # prompts use for data augmentation\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 3. Initialize Embedding Generated Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fate_llm.model_zoo.embedding_transformer.st_model import SentenceTransformerModel\n",
    "\n",
    "\n",
    "embedding_lm = SentenceTransformerModel(model_name_or_path=\"all-mpnet-base-v2\").load() # modified model_name_or_path to local model saved path"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 4. Initalize OpenAI Api For Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fate_llm.algo.fdkt.inference_inst import api_init\n",
    "\n",
    "\n",
    "inference_inst = api_init(\n",
    "    api_url=\"http://127.0.0.1:9999/v1/\",\n",
    "    model_name=\"Qwen1.5-7B-Chat\", # modified model_name to local Meta-Llama-3-8B-Instruct saved path\n",
    "    api_key=\"demo\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Complete Code \n",
    "\n",
    "Please paste the code in \"run_fdkt_by_launcher.py\" and execute it with the following command. Once the process is finished, augmentation data will be saved in the current directory, whose filename is aug_data_result.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "python run_fdkt_by_launcher.py --parties guest:9999 arbiter:10000 --log_level INFO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "\n",
    "import torch\n",
    "from fate.arch import Context\n",
    "from fate.arch.launchers.multiprocess_launcher import launch\n",
    "\n",
    "# please replace the following four variables to local paths\n",
    "llm_pretrained_path = \"Qwen1.5-7B-Chat\"\n",
    "embedding_model_path = \"all-mpnet-base-v2\"\n",
    "slm_pretrained_path = \"gpt2-xl\"\n",
    "slm_data_path = \"./processed_data/Health/train.json\"\n",
    "\n",
    "\n",
    "def get_optimizer(model, optimizer=\"adam\", lr=1e-4):\n",
    "    if optimizer == \"adam\":\n",
    "        optimizer = torch.optim.Adam(params=model.parameters(), lr=lr)\n",
    "    elif optimizer == \"adamw\":\n",
    "        optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr)\n",
    "    else:\n",
    "        raise NotImplementedError(\"Given optimizer type is not supported\")\n",
    "    return optimizer\n",
    "\n",
    "\n",
    "def train_slm(ctx):\n",
    "    import transformers\n",
    "    from fate_llm.algo.fdkt.fdkt_data_aug import (\n",
    "        FDKTSLM,\n",
    "        FDKTTrainingArguments\n",
    "    )\n",
    "    from fate_llm.dataset.flex_dataset import FlexDataset\n",
    "    from fate_llm.data.tokenizers.cust_tokenizer import get_tokenizer\n",
    "    from transformers.data import DataCollatorForSeq2Seq\n",
    "\n",
    "    slm = transformers.AutoModelForCausalLM.from_pretrained(slm_pretrained_path, torch_dtype=torch.bfloat16)\n",
    "    tokenizer = get_tokenizer(slm_pretrained_path)\n",
    "    tokenizer.pad_token_id = tokenizer.eos_token_id\n",
    "    training_args = FDKTTrainingArguments(\n",
    "        use_cpu=False,\n",
    "        device_id=0,\n",
    "        num_train_epochs=1,\n",
    "        per_device_train_batch_size=2,\n",
    "        slm_generation_batch_size=32,\n",
    "        seq_num_for_single_category=2000,\n",
    "        slm_generation_config=dict(\n",
    "            max_new_tokens=256,\n",
    "            do_sample=True,\n",
    "            temperature=1.0,\n",
    "            top_k=50,\n",
    "            top_p=0.9,\n",
    "            repetition_penalty=1.0,\n",
    "            pad_token_id=tokenizer.eos_token_id\n",
    "        ),\n",
    "        # inference_method=\"vllm\",\n",
    "    )\n",
    "\n",
    "    ds = FlexDataset(\n",
    "        tokenizer_name_or_path=slm_pretrained_path,\n",
    "        load_from=\"json\",\n",
    "        data_part=\"train\",\n",
    "        dataset_name=\"yelp_review\",\n",
    "        need_preprocess=True,\n",
    "        select_num=2000,  # use 2000 data to train, default is None, using all data\n",
    "    )\n",
    "    ds.load(slm_data_path)\n",
    "\n",
    "    fdkt_runner = FDKTSLM(\n",
    "        ctx=ctx,\n",
    "        model=slm,\n",
    "        training_args=training_args,\n",
    "        tokenizer=tokenizer,\n",
    "        train_set=ds,\n",
    "        optimizer=get_optimizer(slm),\n",
    "        data_collator=DataCollatorForSeq2Seq(tokenizer, label_pad_token_id=tokenizer.pad_token_id)\n",
    "    )\n",
    "\n",
    "    aug_data = fdkt_runner.aug_data()\n",
    "    with open(\"./aug_data_result.json\", \"w\") as fout:\n",
    "        fout.write(json.dumps(aug_data, indent=4))\n",
    "\n",
    "\n",
    "def train_llm(ctx):\n",
    "    from fate_llm.algo.fdkt.fdkt_data_aug import (\n",
    "        FDKTLLM,\n",
    "        FDKTTrainingArguments\n",
    "    )\n",
    "    from fate_llm.model_zoo.embedding_transformer.st_model import SentenceTransformerModel\n",
    "    from fate_llm.dataset.flex_dataset import FlexDataset\n",
    "    from fate_llm.algo.fdkt.inference_inst import api_init, vllm_init\n",
    "\n",
    "    embedding_lm = SentenceTransformerModel(model_name_or_path=embedding_model_path).load()\n",
    "    training_args = FDKTTrainingArguments(\n",
    "        sample_num_per_cluster=4,\n",
    "        filter_prompt_max_length=2**14,\n",
    "        filter_generation_config=dict(\n",
    "            max_tokens=4096,\n",
    "        ),\n",
    "        use_cpu=False,\n",
    "        aug_generation_config=dict(\n",
    "            max_tokens=4096,\n",
    "            temperature=0.8,\n",
    "            top_p=0.9,\n",
    "        ),\n",
    "        aug_prompt_num=20000,\n",
    "    )\n",
    "\n",
    "    ds = FlexDataset(\n",
    "        tokenizer_name_or_path=llm_pretrained_path,\n",
    "        load_from=\"json\",\n",
    "        data_part=\"train\",\n",
    "        dataset_name=\"yelp_review\",\n",
    "        need_preprocess=True,\n",
    "        few_shot_num_per_label=1,\n",
    "    )\n",
    "\n",
    "    inference_inst = api_init(\n",
    "        api_url=\"http://127.0.0.1:9999/v1/\",\n",
    "        model_name=llm_pretrained_path,\n",
    "        api_key=\"demo\"\n",
    "    )\n",
    "\n",
    "    fdkt_runner = FDKTLLM(\n",
    "        ctx=ctx,\n",
    "        embedding_model=embedding_lm,\n",
    "        training_args=training_args,\n",
    "        dataset=ds,\n",
    "        inference_inst=inference_inst,\n",
    "    )\n",
    "\n",
    "    fdkt_runner.aug_data()\n",
    "\n",
    "\n",
    "def run(ctx: Context):\n",
    "    if ctx.is_on_arbiter:\n",
    "        train_llm(ctx)\n",
    "    else:\n",
    "        os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "        train_slm(ctx)\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    launch(run)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running FDKT with Pipeline (Industrial Using)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Please make sure that FATE and FATE-Flow has been deployed, paste the following code to test_fdkt_by_pipeline.py, the execute \"python test_fdkt_by_pipeline.py\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fate_client.pipeline.components.fate.homo_nn import HomoNN, get_config_of_fdkt_runner\n",
    "from fate_client.pipeline.components.fate.nn.algo_params import FDKTTrainingArguments\n",
    "from fate_client.pipeline.components.fate.nn.loader import LLMModelLoader, LLMDatasetLoader, LLMDataFuncLoader\n",
    "from fate_client.pipeline import FateFlowPipeline\n",
    "from fate_client.pipeline.components.fate.reader import Reader\n",
    "from fate_client.pipeline.components.fate.nn.torch import nn, optim\n",
    "\n",
    "\n",
    "guest = '9999'# replace this party id to actual guest party id in your enviroment\n",
    "arbiter = '9999'# replace this party id to actual arbiter party id in your enviroment\n",
    "\n",
    "# please replace the following four variables to local paths\n",
    "llm_pretrained_path = \"Qwen1.5-7B-Chat\"\n",
    "embedding_model_path = \"all-mpnet-base-v2/\"\n",
    "slm_pretrained_path = \"gpt2-xl\"\n",
    "slm_data_path = \"./processed_data/Health/train.json\" # should be absolute path\n",
    "\n",
    "\n",
    "def get_llm_conf():\n",
    "    embedding_model = LLMModelLoader(\n",
    "        \"embedding_transformer.st_model\",\n",
    "        \"SentenceTransformerModel\",\n",
    "        model_name_or_path=embedding_model_path\n",
    "    )\n",
    "\n",
    "    dataset = LLMDatasetLoader(\n",
    "        \"flex_dataset\",\n",
    "        \"FlexDataset\",\n",
    "        tokenizer_name_or_path=llm_pretrained_path,\n",
    "        need_preprocess=True,\n",
    "        dataset_name=\"yelp_review\",\n",
    "        data_part=\"train\",\n",
    "        load_from=\"json\",\n",
    "        few_shot_num_per_label=1,\n",
    "    )\n",
    "\n",
    "    training_args = FDKTTrainingArguments(\n",
    "        sample_num_per_cluster=4,\n",
    "        filter_prompt_max_length=2 ** 14,\n",
    "        filter_generation_config=dict(\n",
    "            max_tokens=4096,\n",
    "        ),\n",
    "        use_cpu=False,\n",
    "        aug_generation_config=dict(\n",
    "            max_tokens=4096,\n",
    "            temperature=0.8,\n",
    "            top_p=0.9,\n",
    "        ),\n",
    "        aug_prompt_num=20000,\n",
    "    )\n",
    "\n",
    "    inference_inst_conf = dict(\n",
    "        module_name=\"fate_llm.algo.fdkt.inference_inst\",\n",
    "        item_name=\"api_init\",\n",
    "        kwargs=dict(\n",
    "            api_url=\"http://127.0.0.1:9999/v1/\",\n",
    "            model_name=llm_pretrained_path,\n",
    "            api_key=\"demo\"\n",
    "        )\n",
    "    )\n",
    "\n",
    "    return get_config_of_fdkt_runner(\n",
    "        training_args=training_args,\n",
    "        embedding_model=embedding_model,\n",
    "        dataset=dataset,\n",
    "        inference_inst_conf=inference_inst_conf,\n",
    "    )\n",
    "\n",
    "\n",
    "def get_slm_conf():\n",
    "    slm_model = LLMModelLoader(\n",
    "        \"hf_model\",\n",
    "        \"HFAutoModelForCausalLM\",\n",
    "        pretrained_model_name_or_path=slm_pretrained_path,\n",
    "        torch_dtype=\"bfloat16\",\n",
    "    )\n",
    "\n",
    "    tokenizer = LLMDataFuncLoader(\n",
    "        \"tokenizers.cust_tokenizer\",\n",
    "        \"get_tokenizer\",\n",
    "        tokenizer_name_or_path=slm_pretrained_path,\n",
    "        pad_token_id=50256\n",
    "    )\n",
    "\n",
    "    training_args = FDKTTrainingArguments(\n",
    "        use_cpu=False,\n",
    "        device_id=1,\n",
    "        num_train_epochs=1,\n",
    "        per_device_train_batch_size=2,\n",
    "        slm_generation_batch_size=32,\n",
    "        seq_num_for_single_category=2000,\n",
    "        slm_generation_config=dict(\n",
    "            max_new_tokens=256,\n",
    "            do_sample=True,\n",
    "            temperature=1.0,\n",
    "            top_k=50,\n",
    "            top_p=0.9,\n",
    "            repetition_penalty=1.0,\n",
    "            pad_token_id=50256\n",
    "        ),\n",
    "    )\n",
    "\n",
    "    dataset = LLMDatasetLoader(\n",
    "        \"flex_dataset\",\n",
    "        \"FlexDataset\",\n",
    "        tokenizer_name_or_path=slm_pretrained_path,\n",
    "        need_preprocess=True,\n",
    "        dataset_name=\"yelp_review\",\n",
    "        data_part=\"train\",\n",
    "        load_from=\"json\",\n",
    "        select_num=2000,\n",
    "        few_shot_num_per_label=1,\n",
    "    )\n",
    "\n",
    "    optimizer = optim.Adam(lr=0.01)\n",
    "\n",
    "    return get_config_of_fdkt_runner(\n",
    "        model=slm_model,\n",
    "        tokenizer=tokenizer,\n",
    "        training_args=training_args,\n",
    "        dataset=dataset,\n",
    "        optimizer=optimizer,\n",
    "        data_collator=LLMDataFuncLoader(\n",
    "            \"data_collator.cust_data_collator\",\n",
    "            \"get_seq2seq_data_collator\",\n",
    "            label_pad_token_id=50256,\n",
    "            tokenizer_name_or_path=slm_pretrained_path,\n",
    "            pad_token_id=50256,\n",
    "        ),\n",
    "    )\n",
    "\n",
    "\n",
    "pipeline = FateFlowPipeline().set_parties(guest=guest, arbiter=arbiter)\n",
    "pipeline.bind_local_path(path=slm_data_path, namespace=\"experiment\", name=\"slm_train\")\n",
    "\n",
    "\n",
    "reader_0 = Reader(\"reader_0\", runtime_parties=dict(guest=guest))\n",
    "reader_0.guest.task_parameters(\n",
    "    namespace=\"experiment\",\n",
    "    name=\"slm_train\"\n",
    ")\n",
    "\n",
    "\n",
    "homo_nn_0 = HomoNN(\n",
    "    'homo_nn_0',\n",
    "    train_data=reader_0.outputs[\"output_data\"],\n",
    "    runner_module=\"fdkt_runner\",\n",
    "    runner_class=\"FDKTRunner\",\n",
    ")\n",
    "\n",
    "homo_nn_0.arbiter.task_parameters(\n",
    "    runner_conf=get_llm_conf()\n",
    ")\n",
    "\n",
    "homo_nn_0.guest.task_parameters(\n",
    "    runner_conf=get_slm_conf()\n",
    ")\n",
    "\n",
    "pipeline.add_tasks([reader_0, homo_nn_0])\n",
    "pipeline.conf.set(\"task\", dict(engine_run={\"cores\": 1}))\n",
    "\n",
    "pipeline.compile()\n",
    "pipeline.fit()"
   ]
  }
 ],
 "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.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}