# 自定义与拓展
## 目录
- [自定义数据集](#自定义数据集)
- [自定义模型](#自定义模型)
- [自定义对话模板](#自定义对话模板)

## 自定义数据集
我们支持三种**自定义数据集**的方法.

1. 【推荐】直接命令行传参的方式，指定`--dataset xxx.json yyy.jsonl zzz.csv`, **更加方便支持自定义数据集**, 支持五种数据集格式（即使用`SmartPreprocessor`，支持的数据集格式见下方）, 支持`dataset_id`和`dataset_path`. 不需要修改`dataset_info.json`文件.
2. 添加数据集到`dataset_info.json`中, 比第一种方式更灵活但繁琐, 支持对数据集使用两种预处理器并指定其参数: `RenameColumnsPreprocessor`, `ConversationsPreprocessor`（默认使用`SmartPreprocessor`）. 支持直接修改swift内置的`dataset_info.json`, 或者通过`--custom_dataset_info xxx.json`的方式传入外置的json文件（方便pip install而非git clone的用户拓展数据集）.
3. **注册数据集**的方式: 比第1、2种方式更加灵活但繁琐, 支持使用函数对数据集进行预处理. 方法1、2在实现上借助了方法3. 可以直接修改源码进行拓展, 或者通过`--custom_register_path xxx.py`的方式传入, 脚本会对py文件进行解析（方便pip install的用户）.

### 📌 【推荐】直接命令行传参
支持直接传入行自定义的**dataset_id**(兼容MS和HF)和**dataset_path**, 以及同时传入多个自定义数据集以及对应采样数, 脚本会进行自动的预处理和拼接. 如果传入的是`dataset_id`, 默认会使用dataset\_id中的'default'子数据集, 并设置split为'train'. 如果该dataset\_id已经注册, 则会使用注册时传入的subsets、split以及预处理函数. 如果传入的是`dataset_path`, 则可以指定为相对路径和绝对路径, 其中相对路径为相对于当前运行目录.

```bash
--dataset {dataset_id} {dataset_path}

# 数据集混合: 以下取dataset_id中subset1和subset2子数据集并采样20000条
--dataset {dataset_name}#20000 {dataset_id}:{subset1}/{subset2}#20000 {dataset_path}#10000
```

脚本支持的文件格式包含`csv`, `json`, `jsonl`格式. 你需要将传入的文件符合以下数据集格式（只列出了一部分）. 以下格式都支持system (需要注意的是, csv如果指定了system字段, 则无法设置为`None`, 只能指定为空字符串. json和jsonl没有这个限制). `json`, `jsonl`格式的文件支持多轮对话 (`csv`不支持).


**格式1:**

预训练:

```csv
response
11111
aaaaa
AAAAA
```

```jsonl
{"response": "11111"}
{"response": "aaaaa"}
{"response": "AAAAA"}
```

单轮对话:

```csv
system,query,response
00000,11111,22222
00001,aaaaa,bbbbb
00002,AAAAA,BBBBB
```

```jsonl
{"system": "00000", "query": "11111", "response": "22222"}
{"query": "aaaaa", "response": "bbbbb"}
{"query": "AAAAA", "response": "BBBBB"}
```

多轮对话:

```jsonl
{"system": "00000", "query": "55555", "response": "66666"}
{"query": "eeeee", "response": "fffff", "history": []}
{"query": "EEEEE", "response": "FFFFF", "history": [["AAAAA", "BBBBB"], ["CCCCC", "DDDDD"]]}
```

```json
[{"system": "00000", "query": "55555", "response": "66666"},
{"query": "eeeee", "response": "fffff", "history": []},
{"query": "EEEEE", "response": "FFFFF", "history": [["AAAAA", "BBBBB"], ["CCCCC", "DDDDD"]]}]
```

**格式2:**

```jsonl
{"conversations": [{"from": "system", "value": "00000"}, {"from": "user", "value": "11111"}, {"from": "assistant", "value": "22222"}]}
{"conversations": [{"from": "user", "value": "aaaaa"}, {"from": "assistant", "value": "bbbbb"}, {"from": "user", "value": "ccccc"}, {"from": "assistant", "value": "ddddd"}]}
{"conversations": [{"from": "user", "value": "AAAAA"}, {"from": "assistant", "value": "BBBBB"}, {"from": "user", "value": "CCCCC"}, {"from": "assistant", "value": "DDDDD"}]}
```

**格式3:**

```jsonl
{"messages": [{"role": "system", "content": "00000"}, {"role": "user", "content": "11111"}, {"role": "assistant", "content": "22222"}]}
{"messages": [{"role": "user", "content": "aaaaa"}, {"role": "assistant", "content": "bbbbb"}, {"role": "user", "content": "ccccc"}, {"role": "assistant", "content": "ddddd"}]}
{"messages": [{"role": "user", "content": "AAAAA"}, {"role": "assistant", "content": "BBBBB"}, {"role": "user", "content": "CCCCC"}, {"role": "assistant", "content": "DDDDD"}]}
```

**格式4:**

```jsonl
{"system": "00000", "conversation": [{"human": "11111", "assistant": "22222"}]}
{"conversation": [{"human": "aaaaa", "assistant": "bbbbb"}]}
{"conversation": [{"human": "AAAAA", "assistant": "BBBBB"}, {"human": "CCCCC", "assistant": "DDDDD"}, {"human": "EEEEE", "assistant": "FFFFF"}]}
```

**格式5:**

```csv
system,instruction,input,output
00000,11111,22222,33333
00001,aaaaa,bbbbb,ccccc
00002,AAAAA,BBBBB,CCCCC
```

**强化学习（DPO/ORPO）**

```jsonl
{"query": "11111", "response": "22222", "rejected_response": "33333", "history": [["AAAAA", "BBBBB"], ["CCCCC", "DDDDD"]]}
{"query": "aaaaa", "response": "bbbbb", "rejected_response": "ccccc", "history": [["AAAAA", "BBBBB"], ["CCCCC", "DDDDD"]]}
{"query": "AAAAA", "response": "BBBBB", "rejected_response": "CCCCC", "history": [["AAAAA", "BBBBB"], ["CCCCC", "DDDDD"]]}
```

### 添加dataset_info.json
可以参考[swift内置的dataset_info.json](https://github.com/modelscope/swift/blob/main/swift/llm/data/dataset_info.json)进行数据集拓展. 你可以直接在内置的dataset_info.json中添加, 也可以通过`--custom_dataset_info 1.json`传入外置的dataset_info.json的路径、json字符串或者字典.

添加dataset\_id:
```python
# MS
# 使用: `--dataset <dataset_name>`
"<dataset_name>": {
    "dataset_id": "xxx/xxx"
}

# HF
# 使用: `--dataset HF::<dataset_name>` 或者 直接使用`USE_HF`环境变量.
"<dataset_name>": {
    "hf_dataset_id": "xxx/xxx"
}
```

添加dataset\_path:
```python
# 可以指定相对路径和绝对路径. 相对路径相对于dataset_info.json文件所在目录.
# 使用: `--dataset <dataset_name>`
"<dataset_name>": {
    "dataset_path": "xxx"
}
```

支持以下参数:
- dataset\_id: 数据集对应的ModelScope的dataset\_id, 默认为`None`. 最简的设置必须指定`dataset_id`、`hf_dataset_id`和`dataset_path`中的一个.
- subsets: 子数据集的名字列表, 默认为`[]`, 即使用'default'子数据集.
- split: 默认为`['train']`, 通常不需要修改.
- hf\_dataset\_id: 数据集对应的HuggingFace的datasset\_id, 默认为`None`.
- dataset\_path: 用于指定数据集的本地路径, e.g. 1.jsonl等, 默认为`None`. 可以传入相对路径和绝对路径. 如果使用相对路径, 则相对于`dataset_info.json`文件所在目录. 如果设置了dataset\_path, 那么dataset\_id, subsets, hf\_dataset\_id参数失效.
- columns: 默认使用的预处理器为`SmartPreprocessor`, 指定此参数则指定为`RenameColumnsPreprocessor`, 你需要rename数据集中的列并转换为上述**格式1**的样式.
- conversations: 指定此参数则指定预处理器为`ConversationsPreprocessor`  ('columns'的优先级高于'conversations').
- remove\_useless\_columns: 指定是否移除无用的列 (包括: 'query', 'response', 'rejected\_response', 'system', 'history', 'images'), 默认为`True`, 通常不需要设置.
- tags: 用于注释数据集, 默认为`[]`, 通常不需要设置.

如果`dataset_info.json`中参数无法满足您的要求, 例如你需要添加自定义的prompt、需要对数据集提前进行清洗或者进行复杂的数据集获取与预处理, 则可以使用注册数据集的方式, 使用函数的方式来进行数据获取与预处理.


### 注册数据集的方式

以下是一个**注册数据集**的案例. 完整的py文件可以查看[custom.py](https://github.com/modelscope/swift/blob/main/examples/pytorch/llm/custom.py), sh脚本可以查看[custom](https://github.com/modelscope/swift/tree/main/examples/pytorch/llm/scripts/custom). 你可以通过指定`--custom_register_path xxx.py`对注册的内容进行解析.

```python
from typing import Optional, Tuple

from datasets import Dataset as HfDataset
from modelscope import MsDataset

from swift.llm import get_dataset, register_dataset, get_dataset_from_repo
from swift.utils import get_logger

logger = get_logger()


class CustomDatasetName:
    stsb_en = 'stsb-en'

def _preprocess_stsb(dataset: HfDataset) -> HfDataset:
    prompt = """Task: Based on the given two sentences, provide a similarity score between 0.0 and 5.0.
Sentence 1: {text1}
Sentence 2: {text2}
Similarity score: """
    query = []
    response = []
    for d in dataset:
        query.append(prompt.format(text1=d['text1'], text2=d['text2']))
        response.append(f"{d['label']:.1f}")
    return HfDataset.from_dict({'query': query, 'response': response})


register_dataset(CustomDatasetName.stsb_en, 'huangjintao/stsb', None, _preprocess_stsb, get_dataset_from_repo)


if __name__ == '__main__':
    # test dataset
    train_dataset, val_dataset = get_dataset([CustomDatasetName.stsb_en],
                                             check_dataset_strategy='warning')
    print(f'train_dataset: {train_dataset}')
    print(f'val_dataset: {val_dataset}')

```

`register_dataset`会在`DATASET_MAPPING`中注册数据集, 该函数的参数含义如下:

- `dataset_name`: 必填项, 表示数据集的名字, 也是数据集的唯一id.
- `dataset_id_or_path`: 必填项. 表示数据集在ModelScope Hub上的`dataset_id`或者本地的`dataset_dir`.
- `subsets`: 数据集的子数据集列表, 默认为`[]`.
- `preprocess_func`: 预处理函数.
- `get_function`: 默认值为`None`. 获取数据集的函数. 如果传入None, 则使用修饰器方案进行数据集注册. 如果传入一个函数, 则使用正常方案进行注册.
  > `get_function`需要返回`HfDataset`或`Tuple[HfDataset, Optional[HfDataset]]`. 如果只返回一个数据集, 则该数据集为train\_dataset. 如果返回两个数据集, 则分别作为train\_dataset和val\_dataset. `get_dataset`函数支持获取多个数据集, 例如:`get_dataset(['dataset1', 'dataset2'])`. 我们会将各个子数据集的训练集和验证集部分分别进行拼接, 最终返回合并后的训练集和验证集.

  > 函数返回的`HfDataset`需要符合一定的规范. 如果你要进行**预训练**, 那么只需要包含`response`字段, 具体可以参考`'tigerbot-law-zh'`数据集. 如果是**指令微调(单轮对话)**的情况下, 需包含`query`, `response`字段, 分别代表指令微调的用户询问和AI助手的回答, 具体可以参考`'alpaca-zh'`数据集. 如果是**多轮对话**, 则需要额外加上`history`字段, 代表对话的历史信息, 具体可以参考`'damo-agent-mini-zh'`数据集. 如果每个数据集样例具有不同的`system`, 则需要额外加上system字段, 具体你也可以参考`'damo-agent-mini-zh'`数据集.

- `**kwargs`: 其他用于注释数据集的参数. 该参数一般不需要设置.


## 自定义模型
以下是一个**自定义模型**的案例. 完整的py文件可以查看[custom.py](https://github.com/modelscope/swift/blob/main/examples/pytorch/llm/custom.py), sh脚本可以查看[custom](https://github.com/modelscope/swift/tree/main/examples/pytorch/llm/scripts/custom). 你可以通过指定`--custom_register_path xxx.py`对注册的内容进行解析.

```python
from typing import Any, Dict

from modelscope import AutoConfig, AutoModelForCausalLM, AutoTokenizer

from torch import dtype as Dtype
from transformers.utils.versions import require_version

from swift.llm import LoRATM, TemplateType, get_model_tokenizer, register_model
from swift.utils import get_logger

logger = get_logger()


class CustomModelType:
    tigerbot_7b = 'tigerbot-7b'
    tigerbot_13b = 'tigerbot-13b'
    tigerbot_13b_chat = 'tigerbot-13b-chat'


class CustomTemplateType:
    tigerbot = 'tigerbot'


@register_model(CustomModelType.tigerbot_7b,
                'TigerResearch/tigerbot-7b-base-v3', LoRATM.llama,
                TemplateType.default_generation)
@register_model(CustomModelType.tigerbot_13b,
                'TigerResearch/tigerbot-13b-base-v2', LoRATM.llama,
                TemplateType.default_generation)
@register_model(CustomModelType.tigerbot_13b_chat,
                'TigerResearch/tigerbot-13b-chat-v4', LoRATM.llama,
                CustomTemplateType.tigerbot)
def get_tigerbot_model_tokenizer(model_dir: str,
                                 torch_dtype: Dtype,
                                 model_kwargs: Dict[str, Any],
                                 load_model: bool = True,
                                 **kwargs):
    use_flash_attn = kwargs.pop('use_flash_attn', False)
    if use_flash_attn:
        require_version('transformers>=4.34')
        logger.info('Setting use_flash_attention_2: True')
        model_kwargs['use_flash_attention_2'] = True
    model_config = AutoConfig.from_pretrained(
        model_dir, trust_remote_code=True)
    model_config.pretraining_tp = 1
    model_config.torch_dtype = torch_dtype
    logger.info(f'model_config: {model_config}')
    tokenizer = AutoTokenizer.from_pretrained(
        model_dir, trust_remote_code=True)
    model = None
    if load_model:
        model = AutoModelForCausalLM.from_pretrained(
            model_dir,
            config=model_config,
            torch_dtype=torch_dtype,
            trust_remote_code=True,
            **model_kwargs)
    return model, tokenizer


if __name__ == '__main__':
    # test model base
    model, tokenizer = get_model_tokenizer(
        CustomModelType.tigerbot_7b, use_flash_attn=False)
    print(model.__class__.__name__)
    # test model chat
    model, tokenizer = get_model_tokenizer(
        CustomModelType.tigerbot_13b_chat, use_flash_attn=False)
    print(model.__class__.__name__)
```

`register_model`会在`MODEL_MAPPING`中注册模型, 该函数的参数含义如下:

- `model_type`: 必填项. 表示模型的名字, 也是唯一的id.
- `model_id_or_path`: 必填项. 表示模型在ModelScope Hub中的`model_id`, 或者是本地的模型目录`model_dir`.
- `lora_target_modules`: 默认为`None`. 表示在sh脚本中指定`--lora_target_modules DEFAULT`或`--lora_target_modules AUTO`或未指定`--lora_target_modules`情况下默认使用的lora_target_modules.
- `template`: 默认为`TemplateType.default`. 表示在sh脚本中指定`--template_type AUTO`或未指定`--template_type`情况下默认使用的对话模板.
- `get_function`: 默认值为`None`. 获取model和tokenizer的函数. 如果传入None, 则使用修饰器方案进行模型注册. 如果传入一个函数, 则使用正常方案进行注册.
- `requires`: 默认为`[]`. 表示模型所需要的区别于其他模型的依赖. 该参数一般不需要设置.
- `torch_dtype`: 默认为`None`. 表示模型所推荐使用的torch_dtype. 该参数一般不需要设置.
- `revision`: 默认为`None`. 用于指定模型的版本号. 如果`model_id_or_path`是本地的模型目录, 则该参数失效. 该参数一般不需要设置.
- `ignore_file_pattern`: 默认为`None`. 表示下载的时候需要忽略的文件名的正则pattern, 该参数会传递给`snapshot_download`. 例如`r'.+\.bin$'`, `r'.+\.savetensors$'`等. 该参数一般不需要设置.
- `**kwargs`: 其他用于注释模型能力的参数. 该参数一般不需要设置.


## 自定义对话模板
以下是一个**自定义模型**的案例. 完整的py文件可以查看[custom.py](https://github.com/modelscope/swift/blob/main/examples/pytorch/llm/custom.py), sh脚本可以查看[custom](https://github.com/modelscope/swift/tree/main/examples/pytorch/llm/scripts/custom).

```python
from swift.llm import (Template, ModelType, dataset_map,
                       get_model_tokenizer, get_template, get_dataset,
                       print_example, register_template, DatasetName)
from swift.utils import get_logger

logger = get_logger()


class CustomTemplateType:
    tigerbot = 'tigerbot'


# Ref: https://github.com/TigerResearch/TigerBot/blob/main/infer.py
register_template(
    CustomTemplateType.tigerbot,
    Template(['{{SYSTEM}}'], ['\n\n### Instruction:\n{{QUERY}}\n\n### Response:\n'], [],
             [['eos_token_id']]))

if __name__ == '__main__':
    # test template
    train_dataset, _ = get_dataset(DatasetName.blossom_math_zh)
    _, tokenizer = get_model_tokenizer(ModelType.qwen_7b_chat, load_model=False)
    template = get_template(CustomTemplateType.tigerbot, tokenizer)
    train_dataset = dataset_map(train_dataset, template.encode)
    print_example(train_dataset[0], tokenizer)
```

`register_template`会在`TEMPLATE_MAPPING`中注册对话模板, 该函数的参数含义如下:

- `template_type`: 必填项, 表示对话模板的名字, 也是template的唯一id.
- `template`: 必填项, 需要传入一个`Template`. 初始化`Template`需要传入以下参数: `prefix`, `prompt`, `chat_sep`, `suffix`, `default_system`.

模板初始化函数会根据这四个内容, 获取完整的chat template. 其中这四个配置内容的含义如下.

- `prefix`: 表示对话模板中的前缀部分, 一般为system部分, 前缀token, bos token等内容. 我们使用`{{SYSTEM}}`作为system的占位符. 如果`{{SYSTEM}}`没有在prefix中存在, 则该Template不支持system, e.g. `damo-agent-mini-zh`数据集.
- `prompt`: 表示对话模板中的一轮对话. 我们使用`{{QUERY}}`作为每轮对话中, human询问部分的占位符, `{{ROUND0}}`则表示本次对话是第几轮的占位符, 从0开始计数, `{{ROUND1}}`从1开始计数. AI助手的回复部分会拼接在`prompt`的后面, 因此我们没有设计其占位符. 我们只会对AI助手的回复部分计算损失.
- `chat_sep`: 如果需要进行多轮对话, `chat_sep`会作为每轮对话之间的分隔符, 例如: 换行等. 如果设置为None, 则该Template不支持多轮对话.
- `suffix`: 作为对话模板的后缀部分, 一般为eos token. 会拼接在最后一轮的对话后面.
- `default_system`: 默认的system.