Feature 2.2.1 examples

examples/fdkt/fdkt.py

@@ -0,0 +1,164 @@
+import yaml
+from fate_client.pipeline.components.fate.homo_nn import HomoNN, get_config_of_fdkt_runner
+from fate_client.pipeline.components.fate.nn.algo_params import FDKTTrainingArguments
+from fate_client.pipeline.components.fate.nn.loader import LLMModelLoader, LLMDatasetLoader, LLMDataFuncLoader
+from fate_client.pipeline import FateFlowPipeline
+from fate_client.pipeline.components.fate.reader import Reader
+from fate_client.pipeline.components.fate.nn.torch import nn, optim
+from typing import Union, Dict
+import argparse
+
+def main(config="../../config.yaml", param: Union[Dict, str] = None, namespace=""):
+    if isinstance(config, str):
+        config = test_utils.load_job_config(config)
+    if isinstance(param, str):
+        param = yaml.safe_load(param)
+    # Load the configuration file
+    parties = config.parties
+    guest = parties.guest[0]
+    arbiter = parties.arbiter[0]
+
+    pipeline = FateFlowPipeline().set_parties(guest=guest, arbiter=arbiter)
+    pipeline.bind_local_path(path=param["datasets"]["slm_data_path"], namespace=param["data"]["guest"]["namespace"], 
+                                                 name=param["data"]["guest"]["name"])
+
+    def get_llm_conf():
+        embedding_model = LLMModelLoader(
+            "embedding_transformer.st_model",
+            "SentenceTransformerModel",
+            model_name_or_path=param['llm']['embedding_model_path']
+        )
+
+        dataset = LLMDatasetLoader(
+            "flex_dataset",
+            "FlexDataset",
+            tokenizer_name_or_path=param['llm']['pretrained_path'],
+            need_preprocess=True,
+            dataset_name="yelp_review",
+            data_part="train.json",
+            load_from="json",
+            few_shot_num_per_label=1,
+        )
+
+        training_args = FDKTTrainingArguments(
+            sample_num_per_cluster=4,
+            filter_prompt_max_length=2 ** 14,
+            filter_generation_config=dict(
+                max_tokens=3000,
+            ),
+            use_cpu=param['slm']['training_args']['use_cpu'],
+            aug_generation_config=dict(
+                max_tokens=3000,
+                temperature=0.8,
+                top_p=0.9,
+            ),
+            aug_prompt_num=200,
+        )
+
+        inference_inst_conf = dict(
+            module_name="fate_llm.algo.fdkt.inference_inst",
+            item_name="api_init",
+            kwargs=dict(
+                api_url=param['client']['api_url'],
+                model_name=param['llm']['pretrained_path'],
+                api_key=param['client']['api_key']
+            )
+        )
+
+        return get_config_of_fdkt_runner(
+            training_args=training_args,
+            embedding_model=embedding_model,
+            dataset=dataset,
+            inference_inst_conf=inference_inst_conf,
+        )
+
+
+    def get_slm_conf():
+        slm_model = LLMModelLoader(
+            "hf_model",
+            "HFAutoModelForCausalLM",
+            pretrained_model_name_or_path=param['slm']['pretrained_path'],
+            torch_dtype="bfloat16",
+        )
+
+        tokenizer = LLMDataFuncLoader(
+            "tokenizers.cust_tokenizer",
+            "get_tokenizer",
+            tokenizer_name_or_path=param['slm']['pretrained_path'],
+            pad_token_id=50256
+        )
+
+        training_args = FDKTTrainingArguments(
+            use_cpu=param['slm']['training_args']['use_cpu'],
+            device_id=1,
+            num_train_epochs=param['slm']['training_args']['num_train_epochs'],
+            per_device_train_batch_size=param['slm']['training_args']['per_device_train_batch_size'],
+            slm_generation_batch_size=param['slm']['training_args']['slm_generation_batch_size'],
+            seq_num_for_single_category=param['slm']['training_args']['seq_num_for_single_category'],
+            slm_generation_config=param['slm']['training_args']['slm_generation_config'],
+        )
+
+        dataset = LLMDatasetLoader(
+            "flex_dataset",
+            "FlexDataset",
+            tokenizer_name_or_path=param['slm']['pretrained_path'],
+            need_preprocess=True,
+            dataset_name="yelp_review",
+            data_part="train",
+            load_from="json",
+            select_num=2000,
+            few_shot_num_per_label=1,
+        )
+
+        optimizer = optim.Adam(lr=0.01)
+
+        return get_config_of_fdkt_runner(
+            model=slm_model,
+            tokenizer=tokenizer,
+            training_args=training_args,
+            dataset=dataset,
+            optimizer=optimizer,
+            data_collator=LLMDataFuncLoader(
+                "data_collator.cust_data_collator",
+                "get_seq2seq_data_collator",
+                label_pad_token_id=50256,
+                tokenizer_name_or_path=param['slm']['pretrained_path'],
+                pad_token_id=50256,
+            ),
+        )
+
+
+
+    reader_0 = Reader("reader_0", runtime_parties=dict(guest=guest))
+    reader_0.guest.task_parameters(
+        namespace=param["data"]["guest"]["namespace"],
+        name=param["data"]["guest"]["name"]
+    )
+
+    homo_nn_0 = HomoNN(
+        'homo_nn_0',
+        train_data=reader_0.outputs["output_data"],
+        runner_module="fdkt_runner",
+        runner_class="FDKTRunner",
+    )
+
+    homo_nn_0.arbiter.task_parameters(
+        runner_conf=get_llm_conf()
+    )
+
+    homo_nn_0.guest.task_parameters(
+        runner_conf=get_slm_conf()
+    )
+
+    pipeline.add_tasks([reader_0, homo_nn_0])
+    pipeline.conf.set("task", dict(engine_run={"cores": 1}))
+
+    pipeline.compile()
+    pipeline.fit()
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser("FDKT JOB")
+    parser.add_argument("-c", "--config", type=str, help="Path to config file", default="./config.yaml")
+    parser.add_argument("-p", "--param", type=str, help="Path to parameter file", default="./fdkt_config.yaml")
+    args = parser.parse_args()
+    main(args.config, args.param)

examples/fdkt/fdkt_config.yaml

@@ -0,0 +1,97 @@
+data:
+  guest:
+    namespace: experiment
+    name: slm_train
+  host:
+    namespace: experiment
+    name: slm_train
+
+
+datasets:
+  slm_data_path: "train.json" # should be absolute path
+
+# Inference initialization configuration
+
+client:
+  api_url: "http://127.0.0.1:9999/v1"
+  api_key: "demo"
+
+# LLM Configuration
+llm:
+  pretrained_path: "Sheared-LLaMa-1.3B" # Please add your model path
+  embedding_model_path: "all-mpnet-base-v2" # Please add your model path
+
+  dataset:
+    tokenizer_name_or_path: "Sheared-LLaMa-1.3B" # Please add your model path
+    need_preprocess: true
+    dataset_name: "yelp_review"
+    data_part: "train.json"
+    load_from: "json"
+    few_shot_num_per_label: 1
+
+  training_args:
+    sample_num_per_cluster: 4
+    filter_prompt_max_length: 16384
+    filter_generation_config:
+      max_tokens: 3000
+    use_cpu: false
+    aug_generation_config:
+      max_tokens: 3000
+      temperature: 0.8
+      top_p: 0.9
+    aug_prompt_num: 200
+
+  inference_inst_conf:
+    module_name: "fate_llm.algo.fdkt.inference_inst"
+    item_name: "api_init"
+    kwargs:
+      api_url: "http://127.0.0.1:9999/v1/"
+      model_name: "Sheared-LLaMa-1.3B"
+      api_key: "demo"
+
+# SLM Configuration
+slm:
+  pretrained_path: "gpt2" # Please add your model path
+  data_path: "train.json" # Please add your datasets path
+
+  model:
+    torch_dtype: "bfloat16"
+
+  tokenizer:
+    tokenizer_name_or_path: "gpt2" # Please add your model path
+    pad_token_id: 50256
+
+  training_args:
+    use_cpu: false
+    device_id: 1
+    num_train_epochs: 1
+    per_device_train_batch_size: 2
+    slm_generation_batch_size: 32
+    seq_num_for_single_category: 200
+    slm_generation_config:
+      max_new_tokens: 256
+      do_sample: true
+      temperature: 1.0
+      top_k: 50
+      top_p: 0.9
+      repetition_penalty: 1.0
+      pad_token_id: 50256
+
+  dataset:
+    tokenizer_name_or_path: "gpt2" # Please add your model path
+    need_preprocess: true
+    dataset_name: "yelp_review"
+    data_part: "train"
+    load_from: "json"
+    select_num: 2000
+    few_shot_num_per_label: 1
+
+  optimizer:
+    type: "Adam"
+    params:
+      lr: 0.01
+
+  data_collator:
+    label_pad_token_id: 50256
+    tokenizer_name_or_path: "gpt2" # Please add your model path
+    pad_token_id: 50256