diff --git a/README.md b/README.md index f5b6764..ab12b0e 100644 --- a/README.md +++ b/README.md @@ -24,4 +24,5 @@ Use [FATE-LLM deployment packages](https://github.com/FederatedAI/FATE/wiki/Down ## Quick Start - [Federated ChatGLM-6B Training](./doc/tutorial/ChatGLM-6B_ds.ipynb) -- [GPT-2 Training](./doc/tutorial/GPT2-example.ipynb) \ No newline at end of file +- [GPT-2 Training](./doc/tutorial/GPT2-example.ipynb) +- [Builtin Models](./doc/tutorial/builtin_models.md) \ No newline at end of file diff --git a/doc/tutorial/GPT2-example.ipynb b/doc/tutorial/GPT2-example.ipynb index 1b540ff..09f3ee7 100644 --- a/doc/tutorial/GPT2-example.ipynb +++ b/doc/tutorial/GPT2-example.ipynb @@ -521,11 +521,13 @@ "import os\n", "from pipeline import fate_torch_hook\n", "from pipeline.component import HomoNN\n", + "from pipeline.component.homo_nn import DatasetParam, TrainerParam\n", "from pipeline.backend.pipeline import PipeLine\n", "from pipeline.component import Reader\n", "from pipeline.interface import Data\n", "from transformers import GPT2Config\n", "\n", + "\n", "fate_torch_hook(t)\n", "\n", "\n", diff --git a/doc/tutorial/builtin_models.md b/doc/tutorial/builtin_models.md new file mode 100644 index 0000000..06044e7 --- /dev/null +++ b/doc/tutorial/builtin_models.md @@ -0,0 +1,18 @@ +## Builtin Models +FATE-LLM provide some builtin models, users can use them simply to efficiently train their language models. +To use these models, please read the using tutorial of [ChatGLM-6B Training Guide](./ChatGLM-6B_ds.ipynb) and [GPT2 Training Guide](GPT2-example.ipynb). +After reading the training tutorial above, it's easy to use other models listing in the following tabular by changing `module_name`, `class_name`, `dataset` to `ModuleName`, `ClassName`, `DatasetName` respectively list below. + + + +| Model | ModuleName | ClassName | DataSetName | +| -------------- | ----------------- | --------------------------------- | ---------------- | +| LLaMA-7B | pellm.llama | LLAMAForCausalLM | llama_tokenizer | +| ChatGLM-6B | pellm.chatglm | ChatGLMForConditionalGeneration | glm_tokenizer | +| GPT-2 | pellm.gpt2 | GPT2 | nlp_tokenizer | +| ALBERT | pellm.albert | Albert | nlp_tokenizer | +| BART | pellm.bart | Bart | nlp_tokenizer | +| BERT | pellm.bert | Bert | nlp_tokenizer | +| DeBERTa | pellm.deberta | Deberta | nlp_tokenizer | +| DistilBERT | pellm.distilbert | DistilBert | nlp_tokenizer | +| RoBERTa | pellm.roberta | Roberta | nlp_tokenizer |