Upload 396 files

.gitattributes

@@ -32,3 +32,25 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+fengshen/examples/finetune_taiyi_stable_diffusion/demo_dataset/part_0/00000003.jpg filter=lfs diff=lfs merge=lfs -text
+fengshen/examples/stable_diffusion_chinese_EN/result_examples/cat_eating_guoqiao_noodle.png filter=lfs diff=lfs merge=lfs -text
+fengshen/examples/stable_diffusion_chinese_EN/result_examples/huskiy_wearing_space_suit.png filter=lfs diff=lfs merge=lfs -text
+fengshen/examples/stable_diffusion_chinese_EN/result_examples/xiaoqiao_oil_painting.png filter=lfs diff=lfs merge=lfs -text
+fengshen/examples/stable_diffusion_chinese_EN/result_examples/xiaoqiao_vangogh.png filter=lfs diff=lfs merge=lfs -text
+fengshen/examples/stable_diffusion_chinese/img/日出，海面上英文逗号.png filter=lfs diff=lfs merge=lfs -text
+fengshen/examples/stable_diffusion_chinese/img/日出，海面上英文逗号4k壁纸.png filter=lfs diff=lfs merge=lfs -text
+fengshen/examples/stable_diffusion_chinese/img/日出，海面上英文逗号4k壁纸384.png filter=lfs diff=lfs merge=lfs -text
+fengshen/examples/stable_diffusion_chinese/img/日出，海面上英文逗号4k壁纸复杂.png filter=lfs diff=lfs merge=lfs -text
+fengshen/examples/stable_diffusion_chinese/img/日出，海面上英文逗号4k壁纸高清.png filter=lfs diff=lfs merge=lfs -text
+fengshen/examples/stable_diffusion_chinese/img/日出，海面上英文逗号4k壁纸精细.png filter=lfs diff=lfs merge=lfs -text
+fengshen/examples/stable_diffusion_chinese/img/日出，海面上英文逗号插画.png filter=lfs diff=lfs merge=lfs -text
+fengshen/examples/stable_diffusion_chinese/img/日出，海面上英文逗号水彩.png filter=lfs diff=lfs merge=lfs -text
+fengshen/examples/stable_diffusion_chinese/img/日出，海面上英文逗号素描.png filter=lfs diff=lfs merge=lfs -text
+fengshen/examples/stable_diffusion_chinese/img/日出，海面上英文逗号油画.png filter=lfs diff=lfs merge=lfs -text
+fengshen/examples/stable_diffusion_chinese/img/日出，海面上中文逗号.png filter=lfs diff=lfs merge=lfs -text
+fengshen/examples/stable_diffusion_chinese/img/日出，海面上中文感叹号.png filter=lfs diff=lfs merge=lfs -text
+fengshen/examples/stable_diffusion_chinese/img/日出，海面上中文句号.png filter=lfs diff=lfs merge=lfs -text
+fengshen/examples/stable_diffusion_chinese/img/日出，海面上nega广告.png filter=lfs diff=lfs merge=lfs -text
+fengshen/examples/stable_diffusion_chinese/img/日出，海面上nega广告符号.png filter=lfs diff=lfs merge=lfs -text
+fengshen/examples/stable_diffusion_chinese/img/日出，海面上nega广告符号词汇.png filter=lfs diff=lfs merge=lfs -text
+fengshen/examples/stable_diffusion_dreambooth/duck_result.png filter=lfs diff=lfs merge=lfs -text

fengshen/API/main.py

@@ -0,0 +1,76 @@
+import uvicorn
+import click
+import argparse
+import json
+from importlib import import_module
+from fastapi import FastAPI, WebSocket
+from starlette.middleware.cors import CORSMiddleware
+from utils import user_config, api_logger, setup_logger, RequestDataStructure
+
+# 命令行启动时只输入一个参数，即配置文件的名字，eg: text_classification.json
+# 其余所有配置在该配置文件中设定，不在命令行中指定
+total_parser = argparse.ArgumentParser("API")
+total_parser.add_argument("config_path", type=str)
+args = total_parser.parse_args()
+
+# set up user config
+user_config.setup_config(args)
+
+# set up logger
+setup_logger(api_logger, user_config)
+
+# load pipeline 
+pipeline_class = getattr(import_module('fengshen.pipelines.' + user_config.pipeline_type), 'Pipeline')
+model_settings = user_config.model_settings
+model_args = argparse.Namespace(**model_settings)
+pipeline = pipeline_class(
+    args = model_args,
+    model = user_config.model_name
+    )
+
+
+# initialize app
+app = FastAPI(
+    title = user_config.PROJECT_NAME, 
+    openapi_url = f"{user_config.API_PREFIX_STR}/openapi.json"
+    )
+
+
+# api 
+# TODO 
+# 需要针对不同请求方法做不同判断，目前仅跑通了较通用的POST方法
+# POST方法可以完成大多数 输入文本-返回结果 的请求任务
+if(user_config.API_method == "POST"):
+    @app.post(user_config.API_path, tags = user_config.API_tags)
+    async def fengshen_post(data:RequestDataStructure):
+        # logging
+        api_logger.info(data.input_text)
+
+        input_text = data.input_text
+
+        result = pipeline(input_text)
+
+        return result
+else:
+    print("only support POST method")
+
+
+
+# Set all CORS enabled origins
+if user_config.BACKEND_CORS_ORIGINS:
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins = [str(origin) for origin in user_config.BACKEND_CORS_ORIGINS],
+        allow_credentials = user_config.allow_credentials,
+        allow_methods = user_config.allow_methods,
+        allow_headers = user_config.allow_headers,
+    )
+
+
+if __name__ == '__main__':
+
+    # 启动后可在浏览器打开 host:port/docs 查看接口的具体信息，并可进行简单测试
+    # eg: 127.0.0.1:8990/docs
+    uvicorn.run(app, host = user_config.SERVER_HOST, port = user_config.SERVER_PORT)
+     
+

fengshen/API/text_classification.json

@@ -0,0 +1,46 @@
+{
+    "SERVER": {
+        "SERVER_HOST": "127.0.0.1",
+        "SERVER_PORT": 8990,
+        "SERVER_NAME": "fengshen_demo",
+        "PROJECT_NAME": "fengshen_demo",
+        "API_PREFIX_STR": "/api",
+        
+        "API_method" : "POST",
+        "API_path": "/TextClassification",
+        "API_tags": ["TextClassification"],
+
+        "BACKEND_CORS_ORIGINS": ["*"],
+        "allow_credentials": true,
+        "allow_methods": ["*"],
+        "allow_headers": ["*"]
+        
+    },
+    "LOGGING": {
+        "log_file_path": "",
+        "log_level": "INFO"
+    },
+
+    "PIPELINE": {
+        "pipeline_type": "text_classification",
+        "model_name": "IDEA-CCNL/Erlangshen-Roberta-110M-Similarity",
+        "model_settings": {
+            "device": -1,
+            "texta_name": "sentence",
+            "textb_name": "sentence2",
+            "label_name": "label",
+            "max_length": 512,
+            "return_tensors": "pt",
+            "padding": "longest",
+            "truncation": true,
+            "skip_special_tokens": true,
+            "clean_up_tkenization_spaces": true,
+            
+            "skip_steps": 10,
+            "clip_guidance_scale": 7500,
+            "init_scale": 10
+        }
+    }
+}
+
+    

fengshen/API/utils.py

@@ -0,0 +1,167 @@
+from dataclasses import dataclass, field
+import os
+import json
+import logging
+from argparse import Namespace
+from typing import List, Literal, Optional, Union
+from pydantic import AnyHttpUrl, BaseSettings, HttpUrl, validator, BaseModel
+
+
+CURRENT_DIR_PATH = os.path.dirname(os.path.abspath(__file__))
+
+# request body
+# 使用pydantic对请求中的body数据进行验证
+class RequestDataStructure(BaseModel):
+    input_text: List[str] = [""]
+    uuid: Optional[int]
+    
+    # parameters for text2image model
+    input_image: Optional[str]
+    skip_steps: Optional[int]
+    clip_guidance_scale: Optional[int]
+    init_scale: Optional[int]
+
+# API config
+@dataclass
+class APIConfig:
+
+    # server config
+    SERVER_HOST: AnyHttpUrl = "127.0.0.1"
+    SERVER_PORT: int = 8990
+    SERVER_NAME: str = ""
+    PROJECT_NAME: str = ""
+    API_PREFIX_STR: str = "/api"
+
+    # api config
+    API_method: Literal["POST","GET","PUT","OPTIONS","WEBSOCKET","PATCH","DELETE","TRACE","CONNECT"] = "POST"
+    API_path: str = "/TextClassification"
+    API_tags: List[str] = field(default_factory = lambda: [""])
+    
+    # CORS config
+    BACKEND_CORS_ORIGINS: List[AnyHttpUrl] = field(default_factory = lambda: ["*"])
+    allow_credentials: bool = True
+    allow_methods: List[str] = field(default_factory = lambda: ["*"])
+    allow_headers: List[str] = field(default_factory = lambda: ["*"])
+    
+    # log config
+    log_file_path: str = ""
+    log_level: str = "INFO"
+        
+    # pipeline config
+    pipeline_type: str = ""
+    model_name: str = ""
+
+    # model config
+    # device: int = -1
+    # texta_name: Optional[str] = "sentence"
+    # textb_name: Optional[str] = "sentence2"
+    # label_name: Optional[str] = "label"
+    # max_length: int = 512
+    # return_tensors: str = "pt"
+    # padding: str = "longest"
+    # truncation: bool = True
+    # skip_special_tokens: bool = True
+    # clean_up_tkenization_spaces: bool = True
+
+    # # parameters for text2image model
+    # skip_steps: Optional[int] = 0
+    # clip_guidance_scale: Optional[int] = 0
+    # init_scale: Optional[int] = 0
+
+    def setup_config(self, args:Namespace) -> None:
+        
+        # load config file
+        with open(CURRENT_DIR_PATH + "/" + args.config_path, "r") as jsonfile:
+            config = json.load(jsonfile)
+
+        server_config = config["SERVER"]
+        logging_config = config["LOGGING"]
+        pipeline_config = config["PIPELINE"]
+        
+        # server config
+        self.SERVER_HOST: AnyHttpUrl = server_config["SERVER_HOST"]
+        self.SERVER_PORT: int = server_config["SERVER_PORT"]
+        self.SERVER_NAME: str = server_config["SERVER_NAME"]
+        self.PROJECT_NAME: str = server_config["PROJECT_NAME"]
+        self.API_PREFIX_STR: str = server_config["API_PREFIX_STR"]
+
+        # api config
+        self.API_method: Literal["POST","GET","PUT","OPTIONS","WEBSOCKET","PATCH","DELETE","TRACE","CONNECT"] = server_config["API_method"]
+        self.API_path: str = server_config["API_path"]
+        self.API_tags: List[str] = server_config["API_tags"]
+
+        # CORS config
+        self.BACKEND_CORS_ORIGINS: List[AnyHttpUrl] = server_config["BACKEND_CORS_ORIGINS"]
+        self.allow_credentials: bool = server_config["allow_credentials"]
+        self.allow_methods: List[str] = server_config["allow_methods"]
+        self.allow_headers: List[str] = server_config["allow_headers"]
+        
+        # log config
+        self.log_file_path: str = logging_config["log_file_path"]
+        self.log_level: str = logging_config["log_level"]
+        
+        # pipeline config
+        self.pipeline_type: str = pipeline_config["pipeline_type"]
+        self.model_name: str = pipeline_config["model_name"]
+
+        # general model config
+        self.model_settings: dict = pipeline_config["model_settings"]
+
+        # 由于pipeline本身会解析参数，后续参数可以不要
+        # 直接将model_settings字典转为Namespace后作为pipeline的args参数即可
+
+        # self.device: int = self.model_settings["device"]
+        # self.texta_name: Optional[str] = self.model_settings["texta_name"]
+        # self.textb_name: Optional[str] = self.model_settings["textb_name"]
+        # self.label_name: Optional[str] = self.model_settings["label_name"]
+        # self.max_length: int = self.model_settings["max_length"]
+        # self.return_tensors: str = self.model_settings["return_tensors"]
+        # self.padding: str = self.model_settings["padding"]
+        # self.truncation: bool = self.model_settings["truncation"]
+        # self.skip_special_tokens: bool = self.model_settings["skip_special_tokens"]
+        # self.clean_up_tkenization_spaces: bool = self.model_settings["clean_up_tkenization_spaces"]
+
+        # # specific parameters for text2image model
+        # self.skip_steps: Optional[int] = self.model_settings["skip_steps"]
+        # self.clip_guidance_scale: Optional[int] = self.model_settings["clip_guidance_scale"]
+        # self.init_scale: Optional[int] = self.model_settings["init_scale"]
+        
+
+
+def setup_logger(logger, user_config: APIConfig):
+        
+        # default level: INFO 
+
+        logger.setLevel(getattr(logging, user_config.log_level, "INFO"))
+        ch = logging.StreamHandler()
+        
+        if(user_config.log_file_path == ""):
+            fh = logging.FileHandler(filename = CURRENT_DIR_PATH + "/"  + user_config.SERVER_NAME  + ".log")
+        elif(".log" not in user_config.log_file_path[-5:-1]):
+            fh = logging.FileHandler(filename = user_config.log_file_path + "/" + user_config.SERVER_NAME + ".log")
+        else:
+            fh = logging.FileHandler(filename = user_config.log_file_path)
+
+
+        formatter = logging.Formatter(
+            "%(asctime)s - %(module)s - %(funcName)s - line:%(lineno)d - %(levelname)s - %(message)s"
+        )
+
+        ch.setFormatter(formatter)
+        fh.setFormatter(formatter)
+        logger.addHandler(ch)  # Exporting logs to the screen
+        logger.addHandler(fh)  # Exporting logs to a file
+
+        return logger
+
+user_config = APIConfig()
+api_logger = logging.getLogger()
+
+
+
+
+
+
+
+
+

fengshen/README.md

@@ -0,0 +1,105 @@
+## 最新发布
+
+* \[2022.09.13\] [更新ErLangShen系列DeBERTa预训练代码](https://huggingface.co/IDEA-CCNL/Erlangshen-DeBERTa-v2-97M-Chinese)
+* \[2022.09.13\] [更新RanDeng系列Bart预训练代码](https://huggingface.co/IDEA-CCNL/Randeng-BART-139M)
+* \[2022.09.13\] [更新ErLangShen系列Bert预训练代码](https://huggingface.co/IDEA-CCNL/Erlangshen-MegatronBert-1.3B)
+* \[2022.05.11\] [更新TaiYi系列VIT多模态模型及下游任务示例](https://fengshenbang-doc.readthedocs.io/zh/latest/docs/太乙系列/Taiyi-vit-87M-D.html)
+* \[2022.05.11\] [更新BiGan系列Transformer-XL去噪模型及下游任务示例](https://fengshenbang-doc.readthedocs.io/zh/latest/docs/比干系列/Bigan-Transformer-XL-denoise-1.1B.html)
+* \[2022.05.11\] [更新ErLangShen系列下游任务示例](https://fengshenbang-doc.readthedocs.io/zh/latest/docs/二郎神系列/Erlangshen-Roberta-110M-NLI.html)
+
+# 导航
+
+- [导航](#导航)
+  - [框架简介](#框架简介)
+  - [依赖环境](#依赖环境)
+  - [项目结构](#项目结构)
+  - [设计思路](#设计思路)
+  - [分类下游任务](#分类下游任务)
+
+## 框架简介
+
+FengShen训练框架是封神榜大模型开源计划的重要一环，在大模型的生产和应用中起到至关重要的作用。FengShen可以应用在基于海量数据的预训练以及各种下游任务的finetune中。封神榜专注于NLP大模型开源，然而模型的增大带来不仅仅是训练的问题，在使用上也存在诸多不便。为了解决训练和使用的问题，FengShen参考了目前开源的优秀方案并且重新设计了Pipeline，用户可以根据自己的需求，从封神榜中选取丰富的预训练模型，同时利用FengShen快速微调下游任务。
+
+目前所有实例以及文档可以查看我们的[Wiki](https://fengshenbang-doc.readthedocs.io/zh/latest/index.html)
+所有的模型可以在[Huggingface主页](https://huggingface.co/IDEA-CCNL)找到
+
+通过我们的框架，你可以快速享受到：
+
+1. 比原生torch更强的性能，训练速度提升<font color=#0000FF >**300%**</font>
+2. 支持更大的模型，支持<font color=#0000FF >**百亿级别**</font>内模型训练及微调
+3. 支持<font color=#0000FF >**TB级以上**</font>的数据集，在家用主机上即可享受预训练模型带来的效果提升
+3. 丰富的预训练、下游任务示例，一键开始训练
+4. 适应各种设备环境，支持在CPU、GPU、TPU等不同设备上运行
+5. 集成主流的分布式训练逻辑，无需修改代码即可支持DDP、Zero Optimizer等分布式优化技术
+
+![avartar](../pics/fengshen_pic.png)
+
+## 依赖环境
+
+* Python >= 3.8
+* torch >= 1.8
+* transformers >= 3.2.0
+* pytorch-lightning >= 1.5.10
+
+在Fengshenbang-LM根目录下
+pip install --editable ./
+
+## 项目结构
+
+```
+├── data                        # 支持多种数据处理方式以及数据集
+│   ├── cbart_dataloader
+|   ├── fs_datasets             # 基于transformers datasets的封装，新增中文数据集(开源计划中)
+|   ├── universal_datamodule    # 打通fs_datasets与lightning datamodule，减少重复开发工作量
+│   ├── megatron_dataloader     # 支持基于Megatron实现的TB级别数据集处理、训练
+│   ├── mmap_dataloader         # 通用的Memmap形式的数据加载
+│   └── task_dataloader         # 支持多种下游任务
+├── examples                    # 丰富的示例，从预训练到下游任务，应有尽有。
+├── metric                      # 提供各种metric计算，支持用户自定义metric
+├── losses                      # 同样支持loss自定义，满足定制化需求
+├── tokenizer                   # 支持自定义tokenizer，比如我们使用的SentencePiece训练代码等
+├── models                      # 模型库
+│   ├── auto                    # 支持自动导入对应的模型
+│   ├── bart
+│   ├── longformer
+│   ├── megatron_t5
+│   └── roformer
+└── utils                       # 实用函数
+```
+
+## 设计思路
+
+FengShen框架目前整体基于Pytorch-Lightning & Transformer进行开发，在底层框架上不断开源基于中文的预训练模型，同时提供丰富的examples，每一个封神榜的模型都能找到对应的预训练、下游任务代码。
+
+在FengShen上开发，整体可以按照下面的三个步骤进行：
+
+1. 封装数据处理流程 -> pytorch_lightning.LightningDataModule
+2. 封装模型结构 -> pytorch_lightning.LightningModule
+3. 配置一些插件，比如log_monitor，checkpoint_callback等等。
+
+一个完整的DEMO可以看Randeng-BART系列实例 -> [文档](https://fengshenbang-doc.readthedocs.io/zh/latest/docs/燃灯系列/BART-139M.html) [代码](https://github.com/IDEA-CCNL/Fengshenbang-LM/tree/hf-ds/fengshen/examples/pretrain_bart)
+
+## 分类下游任务
+
+ 在examples/classification目录下，我们提供丰富的分类任务的示例���其中我们提供三个一键式运行的示例。
+
+* demo_classification_afqmc_roberta.sh              使用DDP微调roberta
+* demo_classification_afqmc_roberta_deepspeed.sh    结合deepspeed微调roberta，获得更快的运算速度
+* demo_classification_afqmc_erlangshen_offload.sh   仅需7G显存即可微调我们效果最好的二郎神系列模型
+
+ 上述示例均采用AFQMC的数据集，关于数据集的介绍可以在[这里](https://www.cluebenchmarks.com/introduce.html)找到。
+ 同时我们处理过的数据文件已经放在Huggingface上，点击[这里](https://huggingface.co/datasets/IDEA-CCNL/AFQMC)直达源文件。
+ 仅需要按我们的格式稍微处理一下数据集，即可适配下游不同的分类任务。
+ 在脚本示例中，仅需要修改如下参数即可适配本地文件
+
+ ```
+         --dataset_name IDEA-CCNL/AFQMC \
+
+ -------> 修改为
+
+         --data_dir $DATA_DIR \          # 数据目录
+         --train_data train.json \       # 数据文件
+         --valid_data dev.json \
+         --test_data test.json \
+
+ ```

fengshen/__init__.py

@@ -0,0 +1,19 @@
+# coding=utf-8
+# Copyright 2021 The IDEA Authors. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .models.longformer import LongformerConfig, LongformerModel
+from .models.roformer import RoFormerConfig, RoFormerModel
+from .models.megatron_t5 import T5Config, T5EncoderModel
+from .models.ubert import UbertPipelines, UbertModel

fengshen/cli/fengshen_pipeline.py

@@ -0,0 +1,34 @@
+import sys
+from importlib import import_module
+from datasets import load_dataset
+import argparse
+
+
+def main():
+    if len(sys.argv) < 3:
+        raise Exception(
+            'args len < 3, example: fengshen_pipeline text_classification predict xxxxx')
+    pipeline_name = sys.argv[1]
+    method = sys.argv[2]
+    pipeline_class = getattr(import_module('fengshen.pipelines.' + pipeline_name), 'Pipeline')
+
+    total_parser = argparse.ArgumentParser("FengShen Pipeline")
+    total_parser.add_argument('--model', default='', type=str)
+    total_parser.add_argument('--datasets', default='', type=str)
+    total_parser.add_argument('--text', default='', type=str)
+    total_parser = pipeline_class.add_pipeline_specific_args(total_parser)
+    args = total_parser.parse_args(args=sys.argv[3:])
+    pipeline = pipeline_class(args=args, model=args.model)
+
+    if method == 'predict':
+        print(pipeline(args.text))
+    elif method == 'train':
+        datasets = load_dataset(args.datasets)
+        pipeline.train(datasets)
+    else:
+        raise Exception(
+            'cmd not support, now only support {predict, train}')
+
+
+if __name__ == '__main__':
+    main()

fengshen/data/__init__.py

@@ -0,0 +1 @@
+# coding=utf-8

fengshen/data/bert_dataloader/auto_split.sh

@@ -0,0 +1,10 @@
+files=`find $1 -type f -size +1024M`
+
+for p in $files
+do
+echo "processing $p"
+name=`basename $p .json`
+file=`dirname $p`
+split -a 2 -C 300M $p $file/$name- && ls|grep -E "(-[a-zA-Z]{2})" |xargs -n1 -i{} mv {} {}.json
+rm -f $p
+done

fengshen/data/bert_dataloader/load.py

@@ -0,0 +1,200 @@
+import os
+import re
+from pathlib import Path
+import glob
+from tqdm import tqdm
+from contextlib import ExitStack
+import datasets
+import multiprocessing
+from typing import cast, TextIO
+from itertools import chain
+import json
+from concurrent.futures import ProcessPoolExecutor
+from random import shuffle
+from pytorch_lightning import LightningDataModule
+from typing import Optional
+
+from torch.utils.data import DataLoader
+
+
+# _SPLIT_DATA_PATH = '/data1/datas/wudao_180g_split/test'
+_SPLIT_DATA_PATH = '/data1/datas/wudao_180g_split'
+_CACHE_SPLIT_DATA_PATH = '/data1/datas/wudao_180g_FSData'
+
+# feats = datasets.Features({"text": datasets.Value('string')})
+
+
+class BertDataGenerate(object):
+
+    def __init__(self,
+                 data_files=_SPLIT_DATA_PATH,
+                 save_path=_CACHE_SPLIT_DATA_PATH,
+                 train_test_validation='950,49,1',
+                 num_proc=1,
+                 cache=True):
+        self.data_files = Path(data_files)
+        if save_path:
+            self.save_path = Path(save_path)
+        else:
+            self.save_path = self.file_check(
+                Path(self.data_files.parent, self.data_files.name+'_FSDataset'),
+                'save')
+        self.num_proc = num_proc
+        self.cache = cache
+        self.split_idx = self.split_train_test_validation_index(train_test_validation)
+        if cache:
+            self.cache_path = self.file_check(
+                Path(self.save_path.parent, 'FSDataCache', self.data_files.name), 'cache')
+        else:
+            self.cache_path = None
+
+    @staticmethod
+    def file_check(path, path_type):
+        print(path)
+        if not path.exists():
+            path.mkdir(parents=True)
+        print(f"Since no {path_type} directory is specified, the program will automatically create it in {path} directory.")
+        return str(path)
+
+    @staticmethod
+    def split_train_test_validation_index(train_test_validation):
+        split_idx_ = [int(i) for i in train_test_validation.split(',')]
+        idx_dict = {
+            'train_rate': split_idx_[0]/sum(split_idx_),
+            'test_rate': split_idx_[1]/sum(split_idx_[1:])
+        }
+        return idx_dict
+
+    def process(self, index, path):
+        print('saving dataset shard {}'.format(index))
+
+        ds = (datasets.load_dataset('json', data_files=str(path),
+                                    cache_dir=self.cache_path,
+                                    features=None))
+        # ds = ds.map(self.cut_sent,input_columns='text')
+        # print(d)
+        # print('!!!',ds)
+        ds = ds['train'].train_test_split(train_size=self.split_idx['train_rate'])
+        ds_ = ds['test'].train_test_split(train_size=self.split_idx['test_rate'])
+        ds = datasets.DatasetDict({
+            'train': ds['train'],
+            'test': ds_['train'],
+            'validation': ds_['test']
+        })
+        # print('!!!!',ds)
+        ds.save_to_disk(Path(self.save_path, path.name))
+        return 'saving dataset shard {} done'.format(index)
+
+    def generate_cache_arrow(self) -> None:
+        '''
+        生成HF支持的缓存文件，加速后续的加载
+        '''
+        data_dict_paths = self.data_files.rglob('*')
+        p = ProcessPoolExecutor(max_workers=self.num_proc)
+        res = list()
+
+        for index, path in enumerate(data_dict_paths):
+            res.append(p.submit(self.process, index, path))
+
+        p.shutdown(wait=True)
+        for future in res:
+            print(future.result(), flush=True)
+
+
+def load_dataset(num_proc=4, **kargs):
+    cache_dict_paths = Path(_CACHE_SPLIT_DATA_PATH).glob('*')
+    ds = []
+    res = []
+    p = ProcessPoolExecutor(max_workers=num_proc)
+    for path in cache_dict_paths:
+        res.append(p.submit(datasets.load_from_disk,
+                            str(path), **kargs))
+
+    p.shutdown(wait=True)
+    for future in res:
+        ds.append(future.result())
+        # print(future.result())
+    train = []
+    test = []
+    validation = []
+    for ds_ in ds:
+        train.append(ds_['train'])
+        test.append(ds_['test'])
+        validation.append(ds_['validation'])
+    # ds = datasets.concatenate_datasets(ds)
+    # print(ds)
+    return datasets.DatasetDict({
+        'train': datasets.concatenate_datasets(train),
+        'test': datasets.concatenate_datasets(test),
+        'validation': datasets.concatenate_datasets(validation)
+    })
+
+
+class BertDataModule(LightningDataModule):
+    @ staticmethod
+    def add_data_specific_args(parent_args):
+        parser = parent_args.add_argument_group('Universal DataModule')
+        parser.add_argument('--num_workers', default=8, type=int)
+        parser.add_argument('--train_batchsize', default=32, type=int)
+        parser.add_argument('--val_batchsize', default=32, type=int)
+        parser.add_argument('--test_batchsize', default=32, type=int)
+        parser.add_argument('--datasets_name', type=str)
+        # parser.add_argument('--datasets_name', type=str)
+        parser.add_argument('--train_datasets_field', type=str, default='train')
+        parser.add_argument('--val_datasets_field', type=str, default='validation')
+        parser.add_argument('--test_datasets_field', type=str, default='test')
+        return parent_args
+
+    def __init__(
+        self,
+        tokenizer,
+        collate_fn,
+        args,
+        **kwargs,
+    ):
+        super().__init__()
+        self.datasets = load_dataset(num_proc=args.num_workers)
+        self.tokenizer = tokenizer
+        self.collate_fn = collate_fn
+        self.save_hyperparameters(args)
+
+    def setup(self, stage: Optional[str] = None) -> None:
+        self.train = DataLoader(
+            self.datasets[self.hparams.train_datasets_field],
+            batch_size=self.hparams.train_batchsize,
+            shuffle=True,
+            num_workers=self.hparams.num_workers,
+            collate_fn=self.collate_fn,
+        )
+        self.val = DataLoader(
+            self.datasets[self.hparams.val_datasets_field],
+            batch_size=self.hparams.val_batchsize,
+            shuffle=False,
+            num_workers=self.hparams.num_workers,
+            collate_fn=self.collate_fn,
+        )
+        self.test = DataLoader(
+            self.datasets[self.hparams.test_datasets_field],
+            batch_size=self.hparams.test_batchsize,
+            shuffle=False,
+            num_workers=self.hparams.num_workers,
+            collate_fn=self.collate_fn,
+        )
+        return
+
+    def train_dataloader(self):
+        return self.train
+
+    def val_dataloader(self):
+        return self.val
+
+    def test_dataloader(self):
+        return self.test
+
+
+if __name__ == '__main__':
+    # pre = PreProcessing(_SPLIT_DATA_PATH)
+    # pre.processing()
+
+    dataset = BertDataGenerate(_SPLIT_DATA_PATH, num_proc=16)
+    dataset.generate_cache_arrow()

fengshen/data/bert_dataloader/preprocessing.py

@@ -0,0 +1,110 @@
+import re
+import json
+import multiprocessing
+from tqdm import tqdm
+from pathlib import Path
+from itertools import chain
+
+_SPLIT_DATA_PATH = '/data1/datas/wudao_180g'
+
+
+def cut_sent(path):
+    """
+    中文分句，默认？、。、！、省略号分句，考虑双引号包裹的句子
+    采用分割替换的方式
+    """
+    path = Path(path)
+    # print(path)
+    save_path = str(Path('/data1/datas/wudao_180g_split', path.name))
+    print('处理文件：', save_path)
+    with open(save_path, 'wt', encoding='utf-8') as w:
+        with open(path, 'rt', encoding='utf-8') as f:
+            for para in tqdm(f):
+                para = json.loads(para)
+                para_ = para['text'] + ' '
+                # print('sentence piece......')
+                # pep8中 正则不能些 \? 要写成\\?
+                para_ = re.sub('([？。！\\?\\!…]+)([^”’]|[”’])',
+                               r'\1#####\2', para_)
+                para_ = re.sub('([\\.]{3,})([^”’])', r'\1#####\2', para_)
+
+                # 匹配 \1: 句子结束符紧挨’”  \2: 非句子结束符号，被引号包裹的句子
+                para_ = re.sub(
+                    '([。！？\\?\\!…][”’])([^，。！？\\?\\!]|\\s)', r'\1#####\2', para_)
+                para_ = re.sub(
+                    '([\\.]{3,}[”’])([^，。！？\\?\\!]|\\s)', r'\1#####\2', para_)
+                para_ = re.sub(
+                    '([#]{5})([”’])([^，。！？\\?\\!])', r'\2#####\3', para_)
+                para_ = para_.strip()
+                # 一个512里面多个样本
+                line_ = ''
+                for line in para_.split('#####'):
+                    line = line.strip()
+                    if len(line_) < 512 and len(line) > 0:
+                        line_ += line
+                    else:
+                        w.writelines(json.dumps(
+                            {'text': line_}, ensure_ascii=False)+'\n')
+                        line_ = line
+                w.writelines(json.dumps(
+                    {'text': line_}, ensure_ascii=False)+'\n')
+
+
+def chain_iter(*filenames):
+    """
+    将多个文件读成一个迭代器
+    """
+    reader = [open(file, 'r') for file in filenames]
+    return chain(*reader)
+
+
+class Config(object):
+
+    def __init__(self, data_path=_SPLIT_DATA_PATH, num_worker=16, split_numb=600000, cut_sentence=True, output_file=None) -> None:
+        self.data_path = Path(data_path)
+        self.num_worker = num_worker
+        self.split_numb = split_numb
+        self.cut_sentence = cut_sentence
+
+
+def processing1():
+    args = Config()
+    p_ = [str(i) for i in args.data_path.glob('*')]
+    fin = chain_iter(*p_)
+    pool = multiprocessing.Pool(args.num_worker)
+    docs = pool.imap(cut_sent, fin, chunksize=args.num_worker)
+
+    if not Path(args.data_path.parent, args.data_path.name+'_split').exists():
+        Path(args.data_path.parent, args.data_path.name+'_split').mkdir()
+    writer = open(str(Path(args.data_path.parent, args.data_path.name +
+                  '_split', 'sentence_level.json')), 'wt', encoding='utf-8')
+    for doc in tqdm(docs):
+        for sentence in doc:
+            writer.writelines(json.dumps(
+                {"text": sentence}, ensure_ascii=False)+'\n')
+    pool.close()
+    pool.join()
+    writer.close()
+
+
+if __name__ == '__main__':
+    from time import process_time, perf_counter
+    from random import shuffle
+    st = process_time()
+    args = Config(num_worker=16)
+
+    if not Path(args.data_path.parent, args.data_path.name+'_split').exists():
+        Path(args.data_path.parent, args.data_path.name +
+             '_split').mkdir(parents=True)
+
+    p_ = [str(i) for i in args.data_path.glob('*')]
+    # 简单shuffle
+    shuffle(p_)
+
+    pool = multiprocessing.Pool(args.num_worker)
+    for item in p_:
+        pool.apply_async(func=cut_sent, args=(item,))
+    pool.close()
+    pool.join()
+    cost_time = process_time() - st
+    print('DONE!! cost time : %.5f' % cost_time)

fengshen/data/clip_dataloader/flickr.py

@@ -0,0 +1,105 @@
+from torch.utils.data import Dataset, DataLoader
+from torchvision.transforms import Normalize, Compose, RandomResizedCrop, InterpolationMode, ToTensor, Resize, \
+    CenterCrop
+from transformers import BertTokenizer
+import pytorch_lightning as pl
+from PIL import Image
+import os
+
+
+class flickr30k_CNA(Dataset):
+    def __init__(self, img_root_path,
+                 annot_path,
+                 transform=None):
+        self.images = []
+        self.captions = []
+        self.labels = []
+        self.root = img_root_path
+        with open(annot_path, 'r') as f:
+            for line in f:
+                line = line.strip().split('\t')
+                key, caption = line[0].split('#')[0], line[1]
+                img_path = key + '.jpg'
+                self.images.append(img_path)
+                self.captions.append(caption)
+                self.labels.append(key)
+        self.transforms = transform
+        self.tokenizer = BertTokenizer.from_pretrained("hfl/chinese-roberta-wwm-ext")
+
+        # NOTE large 模型
+        self.context_length = 77
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, idx):
+        img_path = str(self.images[idx])
+        image = self.transforms(Image.open(os.path.join(self.root, img_path)))
+        text = self.tokenizer(str(self.captions[idx]), max_length=self.context_length,
+                              padding='max_length', truncation=True, return_tensors='pt')['input_ids'][0]
+        label = self.labels[idx]
+        return image, text, label
+
+
+def _convert_to_rgb(image):
+    return image.convert('RGB')
+
+
+def image_transform(
+        image_size: int,
+        is_train: bool,
+        mean=(0.48145466, 0.4578275, 0.40821073),
+        std=(0.26862954, 0.26130258, 0.27577711)
+):
+    normalize = Normalize(mean=mean, std=std)
+    if is_train:
+        return Compose([
+            RandomResizedCrop(image_size, scale=(0.9, 1.0), interpolation=InterpolationMode.BICUBIC),
+            _convert_to_rgb,
+            ToTensor(),
+            normalize,
+        ])
+    else:
+        return Compose([
+            Resize(image_size, interpolation=InterpolationMode.BICUBIC),
+            CenterCrop(image_size),
+            _convert_to_rgb,
+            ToTensor(),
+            normalize,
+        ])
+
+
+class FlickrDataModule(pl.LightningDataModule):
+    def __init__(self, args):
+        self.batch_size = args.batch_size
+        self.train_filename = args.train_filename  # NOTE 标注的文件夹
+        self.train_root = args.train_root  # NOTE 图片地址
+        self.val_filename = args.val_filename
+        self.val_root = args.val_root
+        self.test_filename = args.test_filename
+        self.test_root = args.test_root
+
+        self.pretrain_model = args.pretrain_model
+        self.image_size = 224
+        self.prepare_data_per_node = True
+        self._log_hyperparams = False
+        self.num_workers = args.num_workers
+
+    def setup(self, stage=None):
+        # dataset
+        train_transform = image_transform(224, True)
+        val_transform = image_transform(224, False)
+        test_transform = image_transform(224, False)
+
+        self.train_dataset = flickr30k_CNA(self.train_root, self.train_filename, transform=train_transform)
+        self.val_dataset = flickr30k_CNA(self.val_root, self.val_filename, transform=val_transform)
+        self.test_dataset = flickr30k_CNA(self.test_root, self.test_filename, transform=test_transform)
+
+    def train_dataloader(self):
+        return DataLoader(self.train_dataset, batch_size=self.batch_size, num_workers=self.num_workers)
+
+    def val_dataloader(self):
+        return DataLoader(self.val_dataset, batch_size=self.batch_size, num_workers=self.num_workers)
+
+    def test_dataloader(self):
+        return DataLoader(self.test_dataset, batch_size=self.batch_size, num_workers=self.num_workers)

fengshen/data/data_utils/common_utils.py

@@ -0,0 +1,4 @@
+def padding_to_maxlength(ids, max_length, pad_id):
+    cur_len = len(ids)
+    len_diff = max_length - len(ids)
+    return ids + [pad_id] * len_diff, [1] * cur_len + [0] * len_diff

fengshen/data/data_utils/mask_utils.py

@@ -0,0 +1,285 @@
+import collections
+
+import numpy as np
+
+MaskedLmInstance = collections.namedtuple("MaskedLmInstance",
+                                          ["index", "label"])
+
+
+def is_start_piece(piece):
+    """Check if the current word piece is the starting piece (BERT)."""
+    # When a word has been split into
+    # WordPieces, the first token does not have any marker and any subsequence
+    # tokens are prefixed with ##. So whenever we see the ## token, we
+    # append it to the previous set of word indexes.
+    return not piece.startswith("##")
+
+
+def create_masked_lm_predictions(tokens,
+                                 vocab_id_list, vocab_id_to_token_dict,
+                                 masked_lm_prob,
+                                 cls_id, sep_id, mask_id,
+                                 max_predictions_per_seq,
+                                 np_rng,
+                                 max_ngrams=3,
+                                 do_whole_word_mask=True,
+                                 favor_longer_ngram=False,
+                                 do_permutation=False,
+                                 geometric_dist=False,
+                                 masking_style="bert",
+                                 zh_tokenizer=None):
+    """Creates the predictions for the masked LM objective.
+    Note: Tokens here are vocab ids and not text tokens."""
+    '''
+    modified from Megatron-LM
+    Args:
+        tokens: 输入
+        vocab_id_list: 词表token_id_list
+        vocab_id_to_token_dict： token_id到token字典
+        masked_lm_prob：mask概率
+        cls_id、sep_id、mask_id：特殊token
+        max_predictions_per_seq：最大mask个数
+        np_rng：mask随机数
+        max_ngrams：最大词长度
+        do_whole_word_mask：是否做全词掩码
+        favor_longer_ngram：优先用长的词
+        do_permutation：是否打乱
+        geometric_dist：用np_rng.geometric做随机
+        masking_style：mask类型
+        zh_tokenizer：WWM的分词器，比如用jieba.lcut做分词之类的
+    '''
+    cand_indexes = []
+    # Note(mingdachen): We create a list for recording if the piece is
+    # the starting piece of current token, where 1 means true, so that
+    # on-the-fly whole word masking is possible.
+    token_boundary = [0] * len(tokens)
+    # 如果没有指定中文分词器，那就直接按##算
+    if zh_tokenizer is None:
+        for (i, token) in enumerate(tokens):
+            if token == cls_id or token == sep_id:
+                token_boundary[i] = 1
+                continue
+        # Whole Word Masking means that if we mask all of the wordpieces
+        # corresponding to an original word.
+        #
+        # Note that Whole Word Masking does *not* change the training code
+        # at all -- we still predict each WordPiece independently, softmaxed
+        # over the entire vocabulary.
+            if (do_whole_word_mask and len(cand_indexes) >= 1 and
+                    not is_start_piece(vocab_id_to_token_dict[token])):
+                cand_indexes[-1].append(i)
+            else:
+                cand_indexes.append([i])
+                if is_start_piece(vocab_id_to_token_dict[token]):
+                    token_boundary[i] = 1
+    else:
+        # 如果指定了中文分词器，那就先用分词器分词，然后再进行判断
+        # 获取去掉CLS SEP的原始文本
+        raw_tokens = []
+        for t in tokens:
+            if t != cls_id and t != sep_id:
+                raw_tokens.append(t)
+        raw_tokens = [vocab_id_to_token_dict[i] for i in raw_tokens]
+        # 分词然后获取每次字开头的最长词的长度
+        word_list = set(zh_tokenizer(''.join(raw_tokens), HMM=True))
+        word_length_dict = {}
+        for w in word_list:
+            if len(w) < 1:
+                continue
+            if w[0] not in word_length_dict:
+                word_length_dict[w[0]] = len(w)
+            elif word_length_dict[w[0]] < len(w):
+                word_length_dict[w[0]] = len(w)
+        i = 0
+        # 从词表里面检索
+        while i < len(tokens):
+            token_id = tokens[i]
+            token = vocab_id_to_token_dict[token_id]
+            if len(token) == 0 or token_id == cls_id or token_id == sep_id:
+                token_boundary[i] = 1
+                i += 1
+                continue
+            word_max_length = 1
+            if token[0] in word_length_dict:
+                word_max_length = word_length_dict[token[0]]
+            j = 0
+            word = ''
+            word_end = i+1
+            # 兼容以前##的形式，如果后面的词是##开头的，那么直接把后面的拼到前面当作一个词
+            old_style = False
+            while word_end < len(tokens) and vocab_id_to_token_dict[tokens[word_end]].startswith('##'):
+                old_style = True
+                word_end += 1
+            if not old_style:
+                while j < word_max_length and i+j < len(tokens):
+                    cur_token = tokens[i+j]
+                    word += vocab_id_to_token_dict[cur_token]
+                    j += 1
+                    if word in word_list:
+                        word_end = i+j
+            cand_indexes.append([p for p in range(i, word_end)])
+            token_boundary[i] = 1
+            i = word_end
+
+    output_tokens = list(tokens)
+
+    masked_lm_positions = []
+    masked_lm_labels = []
+
+    if masked_lm_prob == 0:
+        return (output_tokens, masked_lm_positions,
+                masked_lm_labels, token_boundary)
+
+    num_to_predict = min(max_predictions_per_seq,
+                         max(1, int(round(len(tokens) * masked_lm_prob))))
+
+    ngrams = np.arange(1, max_ngrams + 1, dtype=np.int64)
+    if not geometric_dist:
+        # Note(mingdachen):
+        # By default, we set the probilities to favor shorter ngram sequences.
+        pvals = 1. / np.arange(1, max_ngrams + 1)
+        pvals /= pvals.sum(keepdims=True)
+        if favor_longer_ngram:
+            pvals = pvals[::-1]
+    # 获取一个ngram的idx，对于每个word，记录他的ngram的word
+    ngram_indexes = []
+    for idx in range(len(cand_indexes)):
+        ngram_index = []
+        for n in ngrams:
+            ngram_index.append(cand_indexes[idx:idx + n])
+        ngram_indexes.append(ngram_index)
+
+    np_rng.shuffle(ngram_indexes)
+
+    (masked_lms, masked_spans) = ([], [])
+    covered_indexes = set()
+    for cand_index_set in ngram_indexes:
+        if len(masked_lms) >= num_to_predict:
+            break
+        if not cand_index_set:
+            continue
+        # Note(mingdachen):
+        # Skip current piece if they are covered in lm masking or previous ngrams.
+        for index_set in cand_index_set[0]:
+            for index in index_set:
+                if index in covered_indexes:
+                    continue
+
+        if not geometric_dist:
+            n = np_rng.choice(ngrams[:len(cand_index_set)],
+                              p=pvals[:len(cand_index_set)] /
+                              pvals[:len(cand_index_set)].sum(keepdims=True))
+        else:
+            # Sampling "n" from the geometric distribution and clipping it to
+            # the max_ngrams. Using p=0.2 default from the SpanBERT paper
+            # https://arxiv.org/pdf/1907.10529.pdf (Sec 3.1)
+            n = min(np_rng.geometric(0.2), max_ngrams)
+
+        index_set = sum(cand_index_set[n - 1], [])
+        n -= 1
+        # Note(mingdachen):
+        # Repeatedly looking for a candidate that does not exceed the
+        # maximum number of predictions by trying shorter ngrams.
+        while len(masked_lms) + len(index_set) > num_to_predict:
+            if n == 0:
+                break
+            index_set = sum(cand_index_set[n - 1], [])
+            n -= 1
+        # If adding a whole-word mask would exceed the maximum number of
+        # predictions, then just skip this candidate.
+        if len(masked_lms) + len(index_set) > num_to_predict:
+            continue
+        is_any_index_covered = False
+        for index in index_set:
+            if index in covered_indexes:
+                is_any_index_covered = True
+                break
+        if is_any_index_covered:
+            continue
+        for index in index_set:
+            covered_indexes.add(index)
+            masked_token = None
+            token_id = tokens[index]
+            if masking_style == "bert":
+                # 80% of the time, replace with [MASK]
+                if np_rng.random() < 0.8:
+                    masked_token = mask_id
+                else:
+                    # 10% of the time, keep original
+                    if np_rng.random() < 0.5:
+                        masked_token = tokens[index]
+                    # 10% of the time, replace with random word
+                    else:
+                        masked_token = vocab_id_list[np_rng.randint(0, len(vocab_id_list))]
+            elif masking_style == "t5":
+                masked_token = mask_id
+            else:
+                raise ValueError("invalid value of masking style")
+
+            output_tokens[index] = masked_token
+            masked_lms.append(MaskedLmInstance(index=index, label=token_id))
+
+        masked_spans.append(MaskedLmInstance(
+            index=index_set,
+            label=[tokens[index] for index in index_set]))
+
+    assert len(masked_lms) <= num_to_predict
+    np_rng.shuffle(ngram_indexes)
+
+    select_indexes = set()
+    if do_permutation:
+        for cand_index_set in ngram_indexes:
+            if len(select_indexes) >= num_to_predict:
+                break
+            if not cand_index_set:
+                continue
+            # Note(mingdachen):
+            # Skip current piece if they are covered in lm masking or previous ngrams.
+            for index_set in cand_index_set[0]:
+                for index in index_set:
+                    if index in covered_indexes or index in select_indexes:
+                        continue
+
+            n = np.random.choice(ngrams[:len(cand_index_set)],
+                                 p=pvals[:len(cand_index_set)] /
+                                 pvals[:len(cand_index_set)].sum(keepdims=True))
+            index_set = sum(cand_index_set[n - 1], [])
+            n -= 1
+
+            while len(select_indexes) + len(index_set) > num_to_predict:
+                if n == 0:
+                    break
+                index_set = sum(cand_index_set[n - 1], [])
+                n -= 1
+            # If adding a whole-word mask would exceed the maximum number of
+            # predictions, then just skip this candidate.
+            if len(select_indexes) + len(index_set) > num_to_predict:
+                continue
+            is_any_index_covered = False
+            for index in index_set:
+                if index in covered_indexes or index in select_indexes:
+                    is_any_index_covered = True
+                    break
+            if is_any_index_covered:
+                continue
+            for index in index_set:
+                select_indexes.add(index)
+        assert len(select_indexes) <= num_to_predict
+
+        select_indexes = sorted(select_indexes)
+        permute_indexes = list(select_indexes)
+        np_rng.shuffle(permute_indexes)
+        orig_token = list(output_tokens)
+
+        for src_i, tgt_i in zip(select_indexes, permute_indexes):
+            output_tokens[src_i] = orig_token[tgt_i]
+            masked_lms.append(MaskedLmInstance(index=src_i, label=orig_token[src_i]))
+
+    masked_lms = sorted(masked_lms, key=lambda x: x.index)
+    # Sort the spans by the index of the first span
+    masked_spans = sorted(masked_spans, key=lambda x: x.index[0])
+
+    for p in masked_lms:
+        masked_lm_positions.append(p.index)
+        masked_lm_labels.append(p.label)
+    return (output_tokens, masked_lm_positions, masked_lm_labels, token_boundary, masked_spans)

fengshen/data/data_utils/sentence_split.py

@@ -0,0 +1,35 @@
+import re
+
+
+class ChineseSentenceSplitter(object):
+    def merge_symmetry(self, sentences, symmetry=('“', '”')):
+        # '''合并对称符号，如双引号'''
+        effective_ = []
+        merged = True
+        for index in range(len(sentences)):
+            if symmetry[0] in sentences[index] and symmetry[1] not in sentences[index]:
+                merged = False
+                effective_.append(sentences[index])
+            elif symmetry[1] in sentences[index] and not merged:
+                merged = True
+                effective_[-1] += sentences[index]
+            elif symmetry[0] not in sentences[index] and symmetry[1] not in sentences[index] and not merged:
+                effective_[-1] += sentences[index]
+            else:
+                effective_.append(sentences[index])
+        return [i.strip() for i in effective_ if len(i.strip()) > 0]
+
+    def to_sentences(self, paragraph):
+        #  """由段落切分成句子"""
+        sentences = re.split(r"(？|。|[！]+|!|\…\…)", paragraph)
+        sentences.append("")
+        sentences = ["".join(i) for i in zip(sentences[0::2], sentences[1::2])]
+        sentences = [i.strip() for i in sentences if len(i.strip()) > 0]
+        for j in range(1, len(sentences)):
+            if sentences[j][0] == '”':
+                sentences[j-1] = sentences[j-1] + '”'
+                sentences[j] = sentences[j][1:]
+        return self.merge_symmetry(sentences)
+
+    def tokenize(self, text):
+        return self.to_sentences(text)

fengshen/data/data_utils/sop_utils.py

@@ -0,0 +1,32 @@
+
+# copy from megatron
+def get_a_and_b_segments(sample, np_rng):
+    """Divide sample into a and b segments."""
+
+    # Number of sentences in the sample.
+    n_sentences = len(sample)
+    # Make sure we always have two sentences.
+    assert n_sentences > 1, 'make sure each sample has at least two sentences.'
+
+    # First part:
+    # `a_end` is how many sentences go into the `A`.
+    a_end = 1
+    if n_sentences >= 3:
+        # Note that randin in numpy is exclusive.
+        a_end = np_rng.randint(1, n_sentences)
+    tokens_a = []
+    for j in range(a_end):
+        tokens_a.extend(sample[j])
+
+    # Second part:
+    tokens_b = []
+    for j in range(a_end, n_sentences):
+        tokens_b.extend(sample[j])
+
+    # Random next:
+    is_next_random = False
+    if np_rng.random() < 0.5:
+        is_next_random = True
+        tokens_a, tokens_b = tokens_b, tokens_a
+
+    return tokens_a, tokens_b, is_next_random

fengshen/data/data_utils/token_type_utils.py

@@ -0,0 +1,25 @@
+def create_tokens_and_tokentypes(tokens_a, tokens_b, cls_id, sep_id):
+    """Merge segments A and B, add [CLS] and [SEP] and build tokentypes."""
+
+    tokens = []
+    tokentypes = []
+    # [CLS].
+    tokens.append(cls_id)
+    tokentypes.append(0)
+    # Segment A.
+    for token in tokens_a:
+        tokens.append(token)
+        tokentypes.append(0)
+    # [SEP].
+    tokens.append(sep_id)
+    tokentypes.append(0)
+    # Segment B.
+    for token in tokens_b:
+        tokens.append(token)
+        tokentypes.append(1)
+    if tokens_b:
+        # [SEP].
+        tokens.append(sep_id)
+        tokentypes.append(1)
+
+    return tokens, tokentypes

fengshen/data/data_utils/truncate_utils.py

@@ -0,0 +1,19 @@
+
+def truncate_segments(tokens_a, tokens_b, len_a, len_b, max_num_tokens, np_rng):
+    """Truncates a pair of sequences to a maximum sequence length."""
+    # print(len_a, len_b, max_num_tokens)
+    assert len_a > 0
+    if len_a + len_b <= max_num_tokens:
+        return False
+    while len_a + len_b > max_num_tokens:
+        if len_a > len_b:
+            len_a -= 1
+            tokens = tokens_a
+        else:
+            len_b -= 1
+            tokens = tokens_b
+        if np_rng.random() < 0.5:
+            del tokens[0]
+        else:
+            tokens.pop()
+    return True

fengshen/data/dreambooth_datasets/dreambooth_datasets.py

@@ -0,0 +1,183 @@
+# -*- encoding: utf-8 -*-
+'''
+Copyright 2022 The International Digital Economy Academy (IDEA). CCNL team. All rights reserved.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+@File    :   dreambooth_datasets.py
+@Time    :   2022/11/10 00:20
+@Author  :   Gan Ruyi
+@Version :   1.0
+@Contact :   ganruyi@idea.edu.cn
+@License :   (C)Copyright 2022-2023, CCNL-IDEA
+'''
+from torch.utils.data import Dataset
+from torchvision import transforms
+from PIL import Image
+from pathlib import Path
+
+
+def add_data_args(parent_args):
+    parser = parent_args.add_argument_group('taiyi stable diffusion data args')
+    parser.add_argument(
+        "--instance_data_dir",
+        type=str,
+        default=None,
+        required=True,
+        help="A folder containing the training data of instance images.",
+    )
+    parser.add_argument(
+        "--class_data_dir",
+        type=str,
+        default=None,
+        required=False,
+        help="A folder containing the training data of class images.",
+    )
+    parser.add_argument(
+        "--instance_prompt",
+        type=str,
+        default=None,
+        help="The prompt with identifier specifying the instance",
+    )
+    parser.add_argument(
+        "--class_prompt",
+        type=str,
+        default=None,
+        help="The prompt to specify images in the same class as provided instance images.",
+    )
+    parser.add_argument(
+        "--with_prior_preservation",
+        default=False,
+        action="store_true",
+        help="Flag to add prior preservation loss.",
+    )
+    parser.add_argument("--prior_loss_weight", type=float, default=1.0, help="The weight of prior preservation loss.")
+    parser.add_argument(
+        "--num_class_images",
+        type=int,
+        default=100,
+        help=(
+            "Minimal class images for prior preservation loss. If not have enough images, additional images will be"
+            " sampled with class_prompt."
+        ),
+    )
+    parser.add_argument(
+        "--resolution", type=int, default=512,
+        help=(
+            "The resolution for input images, all the images in the train/validation dataset will be resized to this"
+            " resolution"
+        ),
+    )
+    parser.add_argument(
+        "--center_crop", action="store_true", default=False,
+        help="Whether to center crop images before resizing to resolution"
+    )
+    parser.add_argument(
+        "--sample_batch_size", type=int, default=4, help="Batch size (per device) for sampling images."
+    )
+    return parent_args
+
+
+class DreamBoothDataset(Dataset):
+    """
+    A dataset to prepare the instance and class images with the prompts for fine-tuning the model.
+    It pre-processes the images and the tokenizes prompts.
+    """
+
+    def __init__(
+        self,
+        instance_data_dir,
+        instance_prompt,
+        tokenizer,
+        class_data_dir=None,
+        class_prompt=None,
+        size=512,
+        center_crop=False,
+    ):
+        self.size = size
+        self.center_crop = center_crop
+        self.tokenizer = tokenizer
+
+        self.instance_data_dir = Path(instance_data_dir)
+        if not self.instance_data_dir.exists():
+            raise ValueError("Instance images root doesn't exists.")
+
+        self.instance_images_path = list(Path(instance_data_dir).iterdir())
+        print(self.instance_images_path)
+        self.num_instance_images = len(self.instance_images_path)
+        self.instance_prompt = instance_prompt
+        self._length = self.num_instance_images
+
+        if class_data_dir is not None:
+            self.class_data_dir = Path(class_data_dir)
+            self.class_data_dir.mkdir(parents=True, exist_ok=True)
+            self.class_images_path = list(self.class_data_dir.iterdir())
+            self.num_class_images = len(self.class_images_path)
+            self._length = max(self.num_class_images, self.num_instance_images)
+            self.class_prompt = class_prompt
+        else:
+            self.class_data_dir = None
+
+        self.image_transforms = transforms.Compose(
+            [
+                transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR),
+                transforms.CenterCrop(size) if center_crop else transforms.RandomCrop(size),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5], [0.5]),
+            ]
+        )
+
+    def __len__(self):
+        return self._length
+
+    def __getitem__(self, index):
+        example = {}
+        instance_image = Image.open(self.instance_images_path[index % self.num_instance_images])
+        if not instance_image.mode == "RGB":
+            instance_image = instance_image.convert("RGB")
+        example["instance_images"] = self.image_transforms(instance_image)
+        example["instance_prompt_ids"] = self.tokenizer(
+            self.instance_prompt,
+            padding="do_not_pad",
+            truncation=True,
+            max_length=64,
+            # max_length=self.tokenizer.model_max_length,
+        ).input_ids
+
+        if self.class_data_dir:
+            class_image = Image.open(self.class_images_path[index % self.num_class_images])
+            if not class_image.mode == "RGB":
+                class_image = class_image.convert("RGB")
+            example["class_images"] = self.image_transforms(class_image)
+            example["class_prompt_ids"] = self.tokenizer(
+                self.class_prompt,
+                padding="do_not_pad",
+                truncation=True,
+                # max_length=self.tokenizer.model_max_length,
+                max_length=64,
+            ).input_ids
+
+        return example
+
+
+class PromptDataset(Dataset):
+    "A simple dataset to prepare the prompts to generate class images on multiple GPUs."
+
+    def __init__(self, prompt, num_samples):
+        self.prompt = prompt
+        self.num_samples = num_samples
+
+    def __len__(self):
+        return self.num_samples
+
+    def __getitem__(self, index):
+        example = {}
+        example["prompt"] = self.prompt
+        example["index"] = index
+        return example

fengshen/data/hubert/hubert_dataset.py

@@ -0,0 +1,361 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import itertools
+import logging
+import os
+import sys
+from typing import Any, List, Optional, Union
+
+import numpy as np
+
+import torch
+import torch.nn.functional as F
+from fairseq.data import data_utils
+from fairseq.data.fairseq_dataset import FairseqDataset
+
+logger = logging.getLogger(__name__)
+
+
+def add_data_specific_args(parent_args):
+    parser = parent_args.add_argument_group('Hubert Dataset')
+    parser.add_argument('--data', type=str)
+    parser.add_argument('--sample_rate', type=float, default=16000)
+    parser.add_argument('--label_dir', type=str)
+    parser.add_argument('--labels', type=str, nargs='+')
+    parser.add_argument('--label_rate', type=float)
+    parser.add_argument('--max_keep_size', type=int, default=None)
+    parser.add_argument('--min_sample_size', type=int)
+    parser.add_argument('--max_sample_size', type=int)
+    parser.add_argument('--pad_audio', type=bool)
+    parser.add_argument('--normalize', type=bool)
+    parser.add_argument('--random_crop', type=bool)
+    parser.add_argument('--single_target', type=bool, default=False)
+    return parent_args
+
+
+def load_audio(manifest_path, max_keep, min_keep):
+    n_long, n_short = 0, 0
+    names, inds, sizes = [], [], []
+    with open(manifest_path) as f:
+        root = f.readline().strip()
+        for ind, line in enumerate(f):
+            items = line.strip().split("\t")
+            assert len(items) == 2, line
+            sz = int(items[1])
+            if min_keep is not None and sz < min_keep:
+                n_short += 1
+            elif max_keep is not None and sz > max_keep:
+                n_long += 1
+            else:
+                names.append(items[0])
+                inds.append(ind)
+                sizes.append(sz)
+    tot = ind + 1
+    logger.info(
+        (
+            f"max_keep={max_keep}, min_keep={min_keep}, "
+            f"loaded {len(names)}, skipped {n_short} short and {n_long} long, "
+            f"longest-loaded={max(sizes)}, shortest-loaded={min(sizes)}"
+        )
+    )
+    return root, names, inds, tot, sizes
+
+
+def load_label(label_path, inds, tot):
+    with open(label_path) as f:
+        labels = [line.rstrip() for line in f]
+        assert (
+            len(labels) == tot
+        ), f"number of labels does not match ({len(labels)} != {tot})"
+        labels = [labels[i] for i in inds]
+    return labels
+
+
+def load_label_offset(label_path, inds, tot):
+    with open(label_path) as f:
+        code_lengths = [len(line.encode("utf-8")) for line in f]
+        assert (
+            len(code_lengths) == tot
+        ), f"number of labels does not match ({len(code_lengths)} != {tot})"
+        offsets = list(itertools.accumulate([0] + code_lengths))
+        offsets = [(offsets[i], offsets[i + 1]) for i in inds]
+    return offsets
+
+
+def verify_label_lengths(
+    audio_sizes,
+    audio_rate,
+    label_path,
+    label_rate,
+    inds,
+    tot,
+    tol=0.1,  # tolerance in seconds
+):
+    if label_rate < 0:
+        logger.info(f"{label_path} is sequence label. skipped")
+        return
+
+    with open(label_path) as f:
+        lengths = [len(line.rstrip().split()) for line in f]
+        assert len(lengths) == tot
+        lengths = [lengths[i] for i in inds]
+    num_invalid = 0
+    for i, ind in enumerate(inds):
+        dur_from_audio = audio_sizes[i] / audio_rate
+        dur_from_label = lengths[i] / label_rate
+        if abs(dur_from_audio - dur_from_label) > tol:
+            logger.warning(
+                (
+                    f"audio and label duration differ too much "
+                    f"(|{dur_from_audio} - {dur_from_label}| > {tol}) "
+                    f"in line {ind+1} of {label_path}. Check if `label_rate` "
+                    f"is correctly set (currently {label_rate}). "
+                    f"num. of samples = {audio_sizes[i]}; "
+                    f"label length = {lengths[i]}"
+                )
+            )
+            num_invalid += 1
+    if num_invalid > 0:
+        logger.warning(
+            f"total {num_invalid} (audio, label) pairs with mismatched lengths"
+        )
+
+
+class HubertDataset(FairseqDataset):
+    def __init__(
+        self,
+        manifest_path: str,
+        sample_rate: float,
+        label_paths: List[str],
+        label_rates: Union[List[float], float],  # -1 for sequence labels
+        pad_list: List[str],
+        eos_list: List[str],
+        label_processors: Optional[List[Any]] = None,
+        max_keep_sample_size: Optional[int] = None,
+        min_keep_sample_size: Optional[int] = None,
+        max_sample_size: Optional[int] = None,
+        shuffle: bool = True,
+        pad_audio: bool = False,
+        normalize: bool = False,
+        store_labels: bool = True,
+        random_crop: bool = False,
+        single_target: bool = False,
+    ):
+        self.audio_root, self.audio_names, inds, tot, self.sizes = load_audio(
+            manifest_path, max_keep_sample_size, min_keep_sample_size
+        )
+        self.sample_rate = sample_rate
+        self.shuffle = shuffle
+        self.random_crop = random_crop
+
+        self.num_labels = len(label_paths)
+        self.pad_list = pad_list
+        self.eos_list = eos_list
+        self.label_processors = label_processors
+        self.single_target = single_target
+        self.label_rates = (
+            [label_rates for _ in range(len(label_paths))]
+            if isinstance(label_rates, float)
+            else label_rates
+        )
+        self.store_labels = store_labels
+        if store_labels:
+            self.label_list = [load_label(p, inds, tot) for p in label_paths]
+        else:
+            self.label_paths = label_paths
+            self.label_offsets_list = [
+                load_label_offset(p, inds, tot) for p in label_paths
+            ]
+        assert label_processors is None or len(label_processors) == self.num_labels
+        for label_path, label_rate in zip(label_paths, self.label_rates):
+            verify_label_lengths(
+                self.sizes, sample_rate, label_path, label_rate, inds, tot
+            )
+
+        self.max_sample_size = (
+            max_sample_size if max_sample_size is not None else sys.maxsize
+        )
+        self.pad_audio = pad_audio
+        self.normalize = normalize
+        logger.info(
+            f"pad_audio={pad_audio}, random_crop={random_crop}, "
+            f"normalize={normalize}, max_sample_size={self.max_sample_size}"
+        )
+
+    def get_audio(self, index):
+        import soundfile as sf
+
+        wav_path = os.path.join(self.audio_root, self.audio_names[index])
+        wav, cur_sample_rate = sf.read(wav_path)
+        wav = torch.from_numpy(wav).float()
+        wav = self.postprocess(wav, cur_sample_rate)
+        return wav
+
+    def get_label(self, index, label_idx):
+        if self.store_labels:
+            label = self.label_list[label_idx][index]
+        else:
+            with open(self.label_paths[label_idx]) as f:
+                offset_s, offset_e = self.label_offsets_list[label_idx][index]
+                f.seek(offset_s)
+                label = f.read(offset_e - offset_s)
+
+        if self.label_processors is not None:
+            label = self.label_processors[label_idx](label)
+        return label
+
+    def get_labels(self, index):
+        return [self.get_label(index, i) for i in range(self.num_labels)]
+
+    def __getitem__(self, index):
+        wav = self.get_audio(index)
+        labels = self.get_labels(index)
+        return {"id": index, "source": wav, "label_list": labels}
+
+    def __len__(self):
+        return len(self.sizes)
+
+    def crop_to_max_size(self, wav, target_size):
+        size = len(wav)
+        diff = size - target_size
+        if diff <= 0:
+            return wav, 0
+
+        start, end = 0, target_size
+        if self.random_crop:
+            start = np.random.randint(0, diff + 1)
+            end = size - diff + start
+        return wav[start:end], start
+
+    def collater(self, samples):
+        # target = max(sizes) -> random_crop not used
+        # target = max_sample_size -> random_crop used for long
+        samples = [s for s in samples if s["source"] is not None]
+        if len(samples) == 0:
+            return {}
+
+        audios = [s["source"] for s in samples]
+        audio_sizes = [len(s) for s in audios]
+        if self.pad_audio:
+            audio_size = min(max(audio_sizes), self.max_sample_size)
+        else:
+            audio_size = min(min(audio_sizes), self.max_sample_size)
+        collated_audios, padding_mask, audio_starts = self.collater_audio(
+            audios, audio_size
+        )
+
+        targets_by_label = [
+            [s["label_list"][i] for s in samples] for i in range(self.num_labels)
+        ]
+        targets_list, lengths_list, ntokens_list = self.collater_label(
+            targets_by_label, audio_size, audio_starts
+        )
+
+        net_input = {"source": collated_audios, "padding_mask": padding_mask}
+        batch = {
+            "id": torch.LongTensor([s["id"] for s in samples]),
+            "net_input": net_input,
+        }
+
+        if self.single_target:
+            batch["target_lengths"] = lengths_list[0]
+            batch["ntokens"] = ntokens_list[0]
+            batch["target"] = targets_list[0]
+        else:
+            batch["target_lengths_list"] = lengths_list
+            batch["ntokens_list"] = ntokens_list
+            batch["target_list"] = targets_list
+        return batch
+
+    def collater_audio(self, audios, audio_size):
+        collated_audios = audios[0].new_zeros(len(audios), audio_size)
+        padding_mask = (
+            torch.BoolTensor(collated_audios.shape).fill_(False)
+            # if self.pad_audio else None
+        )
+        audio_starts = [0 for _ in audios]
+        for i, audio in enumerate(audios):
+            diff = len(audio) - audio_size
+            if diff == 0:
+                collated_audios[i] = audio
+            elif diff < 0:
+                assert self.pad_audio
+                collated_audios[i] = torch.cat([audio, audio.new_full((-diff,), 0.0)])
+                padding_mask[i, diff:] = True
+            else:
+                collated_audios[i], audio_starts[i] = self.crop_to_max_size(
+                    audio, audio_size
+                )
+        return collated_audios, padding_mask, audio_starts
+
+    def collater_frm_label(self, targets, audio_size, audio_starts, label_rate, pad):
+        assert label_rate > 0
+        s2f = label_rate / self.sample_rate
+        frm_starts = [int(round(s * s2f)) for s in audio_starts]
+        frm_size = int(round(audio_size * s2f))
+        if not self.pad_audio:
+            rem_size = [len(t) - s for t, s in zip(targets, frm_starts)]
+            frm_size = min(frm_size, *rem_size)
+        targets = [t[s: s + frm_size] for t, s in zip(targets, frm_starts)]
+        logger.debug(f"audio_starts={audio_starts}")
+        logger.debug(f"frame_starts={frm_starts}")
+        logger.debug(f"frame_size={frm_size}")
+
+        lengths = torch.LongTensor([len(t) for t in targets])
+        ntokens = lengths.sum().item()
+        targets = data_utils.collate_tokens(targets, pad_idx=pad, left_pad=False)
+        return targets, lengths, ntokens
+
+    def collater_seq_label(self, targets, pad):
+        lengths = torch.LongTensor([len(t) for t in targets])
+        ntokens = lengths.sum().item()
+        targets = data_utils.collate_tokens(targets, pad_idx=pad, left_pad=False)
+        return targets, lengths, ntokens
+
+    def collater_label(self, targets_by_label, audio_size, audio_starts):
+        targets_list, lengths_list, ntokens_list = [], [], []
+        itr = zip(targets_by_label, self.label_rates, self.pad_list)
+        for targets, label_rate, pad in itr:
+            if label_rate == -1.0:
+                targets, lengths, ntokens = self.collater_seq_label(targets, pad)
+            else:
+                targets, lengths, ntokens = self.collater_frm_label(
+                    targets, audio_size, audio_starts, label_rate, pad
+                )
+            targets_list.append(targets)
+            lengths_list.append(lengths)
+            ntokens_list.append(ntokens)
+        return targets_list, lengths_list, ntokens_list
+
+    def num_tokens(self, index):
+        return self.size(index)
+
+    def size(self, index):
+        if self.pad_audio:
+            return self.sizes[index]
+        return min(self.sizes[index], self.max_sample_size)
+
+    def ordered_indices(self):
+        if self.shuffle:
+            order = [np.random.permutation(len(self))]
+        else:
+            order = [np.arange(len(self))]
+
+        order.append(self.sizes)
+        return np.lexsort(order)[::-1]
+
+    def postprocess(self, wav, cur_sample_rate):
+        if wav.dim() == 2:
+            wav = wav.mean(-1)
+        assert wav.dim() == 1, wav.dim()
+
+        if cur_sample_rate != self.sample_rate:
+            raise Exception(f"sr {cur_sample_rate} != {self.sample_rate}")
+
+        if self.normalize:
+            with torch.no_grad():
+                wav = F.layer_norm(wav, wav.shape)
+        return wav

fengshen/data/megatron_dataloader/Makefile

@@ -0,0 +1,9 @@
+CXXFLAGS += -O3 -Wall -shared -std=c++11 -fPIC -fdiagnostics-color
+CPPFLAGS += $(shell python3 -m pybind11 --includes)
+LIBNAME = helpers
+LIBEXT = $(shell python3-config --extension-suffix)
+
+default: $(LIBNAME)$(LIBEXT)
+
+%$(LIBEXT): %.cpp
+	$(CXX) $(CXXFLAGS) $(CPPFLAGS) $< -o $@

fengshen/data/megatron_dataloader/__init__.py

@@ -0,0 +1 @@
+from . import indexed_dataset

fengshen/data/megatron_dataloader/bart_dataset.py

@@ -0,0 +1,443 @@
+"""BART Style dataset. Modified from fairseq."""
+
+import numpy as np
+import torch
+import math
+import re
+
+from fengshen.data.megatron_dataloader.dataset_utils import (
+    get_samples_mapping
+)
+
+
+class BartDataset(torch.utils.data.Dataset):
+    def __init__(self, name, indexed_dataset, data_prefix,
+                 num_epochs, max_num_samples, masked_lm_prob,
+                 max_seq_length, short_seq_prob, seed, tokenizer, zh_tokenizer):
+
+        # Params to store.
+        self.name = name
+        self.seed = seed
+        self.masked_lm_prob = masked_lm_prob
+        self.max_seq_length = max_seq_length
+
+        # Dataset.
+        self.indexed_dataset = indexed_dataset
+
+        # Build the samples mapping.
+        self.samples_mapping = get_samples_mapping(self.indexed_dataset,
+                                                   data_prefix,
+                                                   num_epochs,
+                                                   max_num_samples,
+                                                   self.max_seq_length - 3,  # account for added tokens
+                                                   short_seq_prob,
+                                                   self.seed,
+                                                   self.name,
+                                                   False)
+
+        # Vocab stuff.
+        self.vocab_size = tokenizer.vocab_size
+        inv_vocab = {v: k for k, v in tokenizer.vocab.items()}
+        self.vocab_id_list = list(inv_vocab.keys())
+        self.vocab_id_to_token_dict = inv_vocab
+        self.cls_id = tokenizer.cls_token_id
+        self.sep_id = tokenizer.sep_token_id
+        self.mask_id = tokenizer.mask_token_id
+        self.pad_id = tokenizer.pad_token_id
+        self.tokenizer = tokenizer
+
+        seg_tokens = ['。', ';', '；', '!', '！', '?', '？']
+        seg_token_ids = []
+        for t in seg_tokens:
+            if t in tokenizer.vocab:
+                seg_token_ids.append(tokenizer.vocab[t])
+            else:
+                print('seg_token "{}" not in vocab'.format(t))
+        self.seg_token_ids = set(seg_token_ids)
+
+        self.zh_tokenizer = zh_tokenizer
+
+        # Denoising ratios
+        self.permute_sentence_ratio = 1.0
+        self.mask_ratio = masked_lm_prob  # 0.15
+        self.random_ratio = 0.1
+        self.insert_ratio = 0.0
+        self.rotate_ratio = 0.0
+        self.mask_whole_word = 1
+        self.item_transform_func = None
+
+        self.mask_span_distribution = None
+        if False:
+            _lambda = 3  # Poisson lambda
+
+            lambda_to_the_k = 1
+            e_to_the_minus_lambda = math.exp(-_lambda)
+            k_factorial = 1
+            ps = []
+            for k in range(0, 128):
+                ps.append(e_to_the_minus_lambda * lambda_to_the_k / k_factorial)
+                lambda_to_the_k *= _lambda
+                k_factorial *= k + 1
+                if ps[-1] < 0.0000001:
+                    break
+            ps = torch.FloatTensor(ps)
+            self.mask_span_distribution = torch.distributions.Categorical(ps)
+
+    def __len__(self):
+        return self.samples_mapping.shape[0]
+
+    def __getitem__(self, idx):
+        start_idx, end_idx, seq_length = self.samples_mapping[idx]
+        sample = [self.indexed_dataset[i] for i in range(start_idx, end_idx)]
+        # Note that this rng state should be numpy and not python since
+        # python randint is inclusive whereas the numpy one is exclusive.
+        # We % 2**32 since numpy requres the seed to be between 0 and 2**32 - 1
+        np_rng = np.random.RandomState(seed=((self.seed + idx) % 2**32))
+        return self.build_training_sample(sample, self.max_seq_length, np_rng)
+
+    def build_training_sample(self, sample, max_seq_length, np_rng):
+        """Biuld training sample.
+
+        Arguments:
+            sample: A list of sentences in which each sentence is a list token ids.
+            max_seq_length: Desired sequence length.
+            np_rng: Random number genenrator. Note that this rng state should be
+                numpy and not python since python randint is inclusive for
+                the opper bound whereas the numpy one is exclusive.
+        """
+        # permute sentences
+        full_stops = []
+        tokens = [self.cls_id]
+        for sent in sample:
+            for t in sent:
+                token = self.vocab_id_to_token_dict[t]
+                if len(re.findall('##[\u4E00-\u9FA5]', token)) > 0:
+                    # 兼容erlangshen ##的方式做whole word mask
+                    t = self.tokenizer.convert_tokens_to_ids(token[2:])
+                tokens.append(t)
+                if t in self.seg_token_ids:
+                    tokens.append(self.sep_id)
+            if tokens[-1] != self.sep_id:
+                tokens.append(self.sep_id)
+
+        if len(tokens) > max_seq_length:
+            tokens = tokens[:max_seq_length]
+            tokens[-1] = self.sep_id
+        tokens = torch.LongTensor(tokens)
+        full_stops = (tokens == self.sep_id).long()
+        assert (max_seq_length - tokens.shape[0]) >= 0, (tokens.size(), tokens[-1], max_seq_length)
+
+        source, target = tokens, tokens[1:].clone()
+        use_decoder = 1
+        # if torch.rand(1).item() < 0.5:
+        #     use_decoder = 0
+
+        if self.permute_sentence_ratio > 0.0 and use_decoder == 1:
+            source = self.permute_sentences(source, full_stops, self.permute_sentence_ratio)
+
+        if self.mask_ratio > 0.0:
+            replace_length = 1 if use_decoder else -1
+            mask_ratio = self.mask_ratio * 2 if use_decoder else self.mask_ratio
+            source = self.add_whole_word_mask(source, mask_ratio, replace_length)
+
+        if self.insert_ratio > 0.0:
+            raise NotImplementedError
+            source = self.add_insertion_noise(source, self.insert_ratio)
+
+        if self.rotate_ratio > 0.0 and np.random.random() < self.rotate_ratio:
+            raise NotImplementedError
+            source = self.add_rolling_noise(source)
+
+        # there can additional changes to make:
+        if self.item_transform_func is not None:
+            source, target = self.item_transform_func(source, target)
+
+        assert (source >= 0).all()
+        # assert (source[1:-1] >= 1).all()
+        assert (source <= self.vocab_size).all()
+        assert source[0] == self.cls_id
+        assert source[-1] == self.sep_id
+
+        # tokenizer = get_tokenizer()
+        # print(' '.join(tokenizer.tokenizer.convert_ids_to_tokens(source)))
+        # print(tokenizer.detokenize(target))
+        # print(tokenizer.detokenize(source))
+        # print()
+
+        prev_output_tokens = torch.zeros_like(target)
+        prev_output_tokens[0] = self.sep_id  # match the preprocessing in fairseq
+        prev_output_tokens[1:] = target[:-1]
+
+        # src_padding_length = max_seq_length - source.shape[0]
+        # tgt_padding_length = max_seq_length - target.shape[0]
+        # assert src_padding_length >= 0, (source.size(), source[-1], max_seq_length)
+        # assert tgt_padding_length >= 0, (target.size(), target[-1], max_seq_length)
+        source_ = torch.full((max_seq_length,), self.pad_id, dtype=torch.long)
+        source_[:source.shape[0]] = source
+        target_ = torch.full((max_seq_length,), -100, dtype=torch.long)
+        # decoder not need bos in the front
+        target_[:target.shape[0]] = target
+        prev_output_tokens_ = torch.full((max_seq_length,), self.pad_id, dtype=torch.long)
+        prev_output_tokens_[:prev_output_tokens.shape[0]] = prev_output_tokens
+
+        return {
+            "input_ids": source_,
+            "labels": target_,
+            # "decoder_input_ids": prev_output_tokens_,
+            "attention_mask": (source_ != self.pad_id).long()
+        }
+
+    def permute_sentences(self, source, full_stops, p=1.0):
+        # Tokens that are full stops, where the previous token is not
+        sentence_ends = (full_stops[1:] * ~full_stops[:-1]).nonzero(as_tuple=False) + 2
+        result = source.clone()
+
+        num_sentences = sentence_ends.size(0)
+        num_to_permute = math.ceil((num_sentences * 2 * p) / 2.0)
+        substitutions = torch.randperm(num_sentences)[:num_to_permute]
+        ordering = torch.arange(0, num_sentences)
+        ordering[substitutions] = substitutions[torch.randperm(num_to_permute)]
+
+        # Ignore <bos> at start
+        index = 1
+        for i in ordering:
+            sentence = source[(sentence_ends[i - 1] if i > 0 else 1): sentence_ends[i]]
+            result[index: index + sentence.size(0)] = sentence
+            index += sentence.size(0)
+        return result
+
+    def word_starts_en(self, source):
+        if self.mask_whole_word is not None:
+            is_word_start = self.mask_whole_word.gather(0, source)
+        else:
+            is_word_start = torch.ones(source.size())
+        is_word_start[0] = 0
+        is_word_start[-1] = 0
+        return is_word_start
+
+    def word_starts(self, source):
+        if self.mask_whole_word is None:
+            is_word_start = torch.ones(source.size())
+            is_word_start[0] = 0
+            is_word_start[-1] = 0
+            return is_word_start
+        raw_tokens = [self.vocab_id_to_token_dict[i] for i in source.tolist()]
+        words = [raw_tokens[0]] + \
+            self.zh_tokenizer(''.join(raw_tokens[1:-1]), HMM=True) + [raw_tokens[-1]]
+
+        def _is_chinese_char(c):
+            """Checks whether CP is the #codepoint of a CJK character."""
+            # This defines a "chinese character" as anything in the CJK Unicode block:
+            #   https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
+            #
+            # Note that the CJK Unicode block is NOT all Japanese and Korean characters,
+            # despite its name. The modern Korean Hangul alphabet is a different block,
+            # as is Japanese Hiragana and Katakana. Those alphabets are used to write
+            # space-separated words, so they are not treated specially and handled
+            # like the all of the other languages.
+            if len(c) > 1:
+                return all([_is_chinese_char(c_i) for c_i in c])
+            cp = ord(c)
+            if ((cp >= 0x4E00 and cp <= 0x9FFF) or  #
+                (cp >= 0x3400 and cp <= 0x4DBF) or  #
+                (cp >= 0x20000 and cp <= 0x2A6DF) or  #
+                (cp >= 0x2A700 and cp <= 0x2B73F) or  #
+                (cp >= 0x2B740 and cp <= 0x2B81F) or  #
+                (cp >= 0x2B820 and cp <= 0x2CEAF) or
+                (cp >= 0xF900 and cp <= 0xFAFF) or  #
+                    (cp >= 0x2F800 and cp <= 0x2FA1F)):  #
+                return True
+
+            return False
+
+        def align_linear(atokens, btokens):
+            a2c = []
+            c2b = []
+            a2b = []
+            length = 0
+            for tok in atokens:
+                a2c.append([length + i for i in range(len(tok))])
+                length += len(tok)
+            for i, tok in enumerate(btokens):
+                c2b.extend([i for _ in range(len(tok))])
+
+            for i, amap in enumerate(a2c):
+                bmap = [c2b[ci] for ci in amap]
+                a2b.append(list(set(bmap)))
+            return a2b
+
+        raw_to_word_align = align_linear(raw_tokens, words)
+        is_word_start = torch.zeros(source.size())
+        word_starts = []
+        skip_cur_word = True
+        for i in range(1, len(raw_to_word_align)):
+            if raw_to_word_align[i-1] == raw_to_word_align[i]:
+                # not a word start, as they align to the same word
+                if not skip_cur_word and not _is_chinese_char(raw_tokens[i]):
+                    word_starts.pop(-1)
+                    skip_cur_word = True
+                continue
+            else:
+                is_word_start[i] = 1
+                if _is_chinese_char(raw_tokens[i]):
+                    word_starts.append(i)
+                    skip_cur_word = False
+        is_word_start[0] = 0
+        is_word_start[-1] = 0
+        word_starts = torch.tensor(word_starts).long().view(-1, 1)
+        return is_word_start, word_starts
+
+    def add_whole_word_mask(self, source, p, replace_length=1):
+        is_word_start, word_starts = self.word_starts(source)
+        num_to_mask_word = int(math.ceil(word_starts.size(0) * p))
+        num_to_mask_char = int(math.ceil(word_starts.size(0) * p * 0.1))
+        num_to_mask = num_to_mask_word + num_to_mask_char
+        if num_to_mask > word_starts.size(0):
+            word_starts = is_word_start.nonzero(as_tuple=False)
+        num_inserts = 0
+        if num_to_mask == 0:
+            return source
+
+        if self.mask_span_distribution is not None:
+            lengths = self.mask_span_distribution.sample(sample_shape=(num_to_mask,))
+
+            # Make sure we have enough to mask
+            cum_length = torch.cumsum(lengths, 0)
+            while cum_length[-1] < num_to_mask:
+                lengths = torch.cat(
+                    [
+                        lengths,
+                        self.mask_span_distribution.sample(sample_shape=(num_to_mask,)),
+                    ],
+                    dim=0,
+                )
+                cum_length = torch.cumsum(lengths, 0)
+
+            # Trim to masking budget
+            i = 0
+            while cum_length[i] < num_to_mask:
+                i += 1
+            lengths[i] = num_to_mask - (0 if i == 0 else cum_length[i - 1])
+            num_to_mask = i + 1
+            lengths = lengths[:num_to_mask]
+
+            # Handle 0-length mask (inserts) separately
+            lengths = lengths[lengths > 0]
+            num_inserts = num_to_mask - lengths.size(0)
+            num_to_mask -= num_inserts
+            if num_to_mask == 0:
+                return self.add_insertion_noise(source, num_inserts / source.size(0))
+
+            assert (lengths > 0).all()
+        else:
+            lengths = torch.ones((num_to_mask,)).long()
+        assert is_word_start[-1] == 0
+        indices = word_starts[
+            torch.randperm(word_starts.size(0))[:num_to_mask]
+        ].squeeze(1)
+        mask_random = torch.FloatTensor(num_to_mask).uniform_() < self.random_ratio
+        source_length = source.size(0)
+        assert source_length - 1 not in indices
+        to_keep = torch.ones(source_length, dtype=torch.bool)
+        is_word_start[
+            -1
+        ] = 255  # acts as a long length, so spans don't go over the end of doc
+        if replace_length == 0:
+            to_keep[indices] = 0
+        else:
+            # keep index, but replace it with [MASK]
+            # print(source.size(), word_starts.size(), indices.size(), mask_random.size())
+            source[indices] = self.mask_id
+            source[indices[mask_random]] = torch.randint(
+                1, self.vocab_size, size=(mask_random.sum(),)
+            )
+            # sorted_indices = torch.sort(indices)[0]
+            # continue_mask_pos = ((sorted_indices + 1)[:-1] == sorted_indices[1:])
+            # continue_mask_indices = sorted_indices[1:][continue_mask_pos]
+            # to_keep[continue_mask_indices] = 0
+
+        # for char indices, we already masked, the following loop handles word mask
+        indices = indices[:num_to_mask_word]
+        mask_random = mask_random[:num_to_mask_word]
+        if self.mask_span_distribution is not None:
+            assert len(lengths.size()) == 1
+            assert lengths.size() == indices.size()
+            lengths -= 1
+            while indices.size(0) > 0:
+                assert lengths.size() == indices.size()
+                lengths -= is_word_start[indices + 1].long()
+                uncompleted = lengths >= 0
+                indices = indices[uncompleted] + 1
+                mask_random = mask_random[uncompleted]
+                lengths = lengths[uncompleted]
+                if replace_length != -1:
+                    # delete token
+                    to_keep[indices] = 0
+                else:
+                    # keep index, but replace it with [MASK]
+                    source[indices] = self.mask_id
+                    source[indices[mask_random]] = torch.randint(
+                        1, self.vocab_size, size=(mask_random.sum(),)
+                    )
+        else:
+            # A bit faster when all lengths are 1
+            while indices.size(0) > 0:
+                uncompleted = is_word_start[indices + 1] == 0
+                indices = indices[uncompleted] + 1
+                mask_random = mask_random[uncompleted]
+                if replace_length != -1:
+                    # delete token
+                    to_keep[indices] = 0
+                else:
+                    # keep index, but replace it with [MASK]
+                    source[indices] = self.mask_id
+                    source[indices[mask_random]] = torch.randint(
+                        1, self.vocab_size, size=(mask_random.sum(),)
+                    )
+
+                assert source_length - 1 not in indices
+
+        source = source[to_keep]
+
+        if num_inserts > 0:
+            source = self.add_insertion_noise(source, num_inserts / source.size(0))
+
+        return source
+
+    def add_permuted_noise(self, tokens, p):
+        num_words = len(tokens)
+        num_to_permute = math.ceil(((num_words * 2) * p) / 2.0)
+        substitutions = torch.randperm(num_words - 2)[:num_to_permute] + 1
+        tokens[substitutions] = tokens[substitutions[torch.randperm(num_to_permute)]]
+        return tokens
+
+    def add_rolling_noise(self, tokens):
+        offset = np.random.randint(1, max(1, tokens.size(-1) - 1) + 1)
+        tokens = torch.cat(
+            (tokens[0:1], tokens[offset:-1], tokens[1:offset], tokens[-1:]),
+            dim=0,
+        )
+        return tokens
+
+    def add_insertion_noise(self, tokens, p):
+        if p == 0.0:
+            return tokens
+
+        num_tokens = len(tokens)
+        n = int(math.ceil(num_tokens * p))
+
+        noise_indices = torch.randperm(num_tokens + n - 2)[:n] + 1
+        noise_mask = torch.zeros(size=(num_tokens + n,), dtype=torch.bool)
+        noise_mask[noise_indices] = 1
+        result = torch.LongTensor(n + len(tokens)).fill_(-1)
+
+        num_random = int(math.ceil(n * self.random_ratio))
+        result[noise_indices[num_random:]] = self.mask_id
+        result[noise_indices[:num_random]] = torch.randint(
+            low=1, high=self.vocab_size, size=(num_random,)
+        )
+
+        result[~noise_mask] = tokens
+
+        assert (result >= 0).all()
+        return result

fengshen/data/megatron_dataloader/bert_dataset.py

@@ -0,0 +1,196 @@
+# coding=utf-8
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""BERT Style dataset."""
+
+
+import numpy as np
+import torch
+
+from fengshen.data.megatron_dataloader.dataset_utils import (
+    get_samples_mapping,
+    get_a_and_b_segments,
+    create_masked_lm_predictions,
+    create_tokens_and_tokentypes,
+)
+
+
+class BertDataset(torch.utils.data.Dataset):
+
+    def __init__(self, name, indexed_dataset, data_prefix,
+                 num_epochs, max_num_samples, masked_lm_prob,
+                 max_seq_length, short_seq_prob, seed, binary_head, tokenizer, masking_style):
+        # Params to store.
+        self.name = name
+        self.seed = seed
+        self.masked_lm_prob = masked_lm_prob
+        self.max_seq_length = max_seq_length
+        self.short_seq_prob = short_seq_prob
+        self.binary_head = binary_head
+        self.masking_style = masking_style
+
+        # Dataset.
+        self.indexed_dataset = indexed_dataset
+
+        # Build the samples mapping.
+        self.samples_mapping = get_samples_mapping(self.indexed_dataset,
+                                                   data_prefix,
+                                                   num_epochs,
+                                                   max_num_samples,
+                                                   # account for added tokens
+                                                   self.max_seq_length - 3,
+                                                   short_seq_prob,
+                                                   self.seed,
+                                                   self.name,
+                                                   self.binary_head)
+        inv_vocab = {v: k for k, v in tokenizer.vocab.items()}
+        self.vocab_id_list = list(inv_vocab.keys())
+        self.vocab_id_to_token_dict = inv_vocab
+        self.cls_id = tokenizer.cls_token_id
+        self.sep_id = tokenizer.sep_token_id
+        self.mask_id = tokenizer.mask_token_id
+        self.pad_id = tokenizer.pad_token_id
+        self.tokenizer = tokenizer
+
+    def __len__(self):
+        return self.samples_mapping.shape[0]
+
+    def __getitem__(self, idx):
+        start_idx, end_idx, seq_length = self.samples_mapping[idx]
+        sample = [self.indexed_dataset[i] for i in range(start_idx, end_idx)]
+        # Note that this rng state should be numpy and not python since
+        # python randint is inclusive whereas the numpy one is exclusive.
+        # We % 2**32 since numpy requres the seed to be between 0 and 2**32 - 1
+        np_rng = np.random.RandomState(seed=((self.seed + idx) % 2**32))
+        return build_training_sample(sample, seq_length,
+                                     self.max_seq_length,  # needed for padding
+                                     self.vocab_id_list,
+                                     self.vocab_id_to_token_dict,
+                                     self.cls_id, self.sep_id,
+                                     self.mask_id, self.pad_id,
+                                     self.masked_lm_prob, np_rng,
+                                     self.binary_head,
+                                     tokenizer=self.tokenizer,
+                                     masking_style=self.masking_style)
+
+
+def build_training_sample(sample,
+                          target_seq_length, max_seq_length,
+                          vocab_id_list, vocab_id_to_token_dict,
+                          cls_id, sep_id, mask_id, pad_id,
+                          masked_lm_prob, np_rng, binary_head,
+                          tokenizer,
+                          masking_style='bert'):
+    """Biuld training sample.
+
+    Arguments:
+        sample: A list of sentences in which each sentence is a list token ids.
+        target_seq_length: Desired sequence length.
+        max_seq_length: Maximum length of the sequence. All values are padded to
+            this length.
+        vocab_id_list: List of vocabulary ids. Used to pick a random id.
+        vocab_id_to_token_dict: A dictionary from vocab ids to text tokens.
+        cls_id: Start of example id.
+        sep_id: Separator id.
+        mask_id: Mask token id.
+        pad_id: Padding token id.
+        masked_lm_prob: Probability to mask tokens.
+        np_rng: Random number genenrator. Note that this rng state should be
+              numpy and not python since python randint is inclusive for
+              the opper bound whereas the numpy one is exclusive.
+    """
+
+    if binary_head:
+        # We assume that we have at least two sentences in the sample
+        assert len(sample) > 1
+    assert target_seq_length <= max_seq_length
+
+    # Divide sample into two segments (A and B).
+    if binary_head:
+        tokens_a, tokens_b, is_next_random = get_a_and_b_segments(sample,
+                                                                  np_rng)
+    else:
+        tokens_a = []
+        for j in range(len(sample)):
+            tokens_a.extend(sample[j])
+        tokens_b = []
+        is_next_random = False
+
+    if len(tokens_a) >= max_seq_length-3:
+        tokens_a = tokens_a[:max_seq_length-3]
+
+    # Truncate to `target_sequence_length`.
+    max_num_tokens = target_seq_length
+    ''''
+    truncated = truncate_segments(tokens_a, tokens_b, len(tokens_a),
+                                  len(tokens_b), max_num_tokens, np_rng)
+    '''
+
+    # Build tokens and toketypes.
+    tokens, tokentypes = create_tokens_and_tokentypes(tokens_a, tokens_b,
+                                                      cls_id, sep_id)
+    # Masking.
+    max_predictions_per_seq = masked_lm_prob * max_num_tokens
+    (tokens, masked_positions, masked_labels, _, _) = create_masked_lm_predictions(
+        tokens, vocab_id_list, vocab_id_to_token_dict, masked_lm_prob,
+        cls_id, sep_id, mask_id, max_predictions_per_seq, np_rng,
+        tokenizer=tokenizer,
+        masking_style=masking_style)
+
+    # Padding.
+    tokens_np, tokentypes_np, labels_np, padding_mask_np, loss_mask_np \
+        = pad_and_convert_to_numpy(tokens, tokentypes, masked_positions,
+                                   masked_labels, pad_id, max_seq_length)
+
+    train_sample = {
+        'input_ids': tokens_np,
+        'token_type_ids': tokentypes_np,
+        'labels': labels_np,
+        'next_sentence_label': int(is_next_random),
+        'attention_mask': padding_mask_np}
+    return train_sample
+
+
+def pad_and_convert_to_numpy(tokens, tokentypes, masked_positions,
+                             masked_labels, pad_id, max_seq_length):
+    """Pad sequences and convert them to numpy."""
+
+    # Some checks.
+    num_tokens = len(tokens)
+    padding_length = max_seq_length - num_tokens
+    assert padding_length >= 0
+    assert len(tokentypes) == num_tokens
+    assert len(masked_positions) == len(masked_labels)
+
+    # Tokens and token types.
+    filler = [pad_id] * padding_length
+    tokens_np = np.array(tokens + filler, dtype=np.int64)
+    tokentypes_np = np.array(tokentypes + filler, dtype=np.int64)
+
+    # Padding mask.
+    padding_mask_np = np.array([1] * num_tokens + [0] * padding_length,
+                               dtype=np.int64)
+
+    # Lables and loss mask.
+    labels = [-100] * max_seq_length
+    loss_mask = [0] * max_seq_length
+    for i in range(len(masked_positions)):
+        assert masked_positions[i] < num_tokens
+        labels[masked_positions[i]] = masked_labels[i]
+        loss_mask[masked_positions[i]] = 1
+    labels_np = np.array(labels, dtype=np.int64)
+    loss_mask_np = np.array(loss_mask, dtype=np.int64)
+
+    return tokens_np, tokentypes_np, labels_np, padding_mask_np, loss_mask_np

fengshen/data/megatron_dataloader/blendable_dataset.py

@@ -0,0 +1,64 @@
+# coding=utf-8
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Blendable dataset."""
+
+import time
+
+import numpy as np
+import torch
+
+from fengshen.data.megatron_dataloader.utils import print_rank_0
+
+
+class BlendableDataset(torch.utils.data.Dataset):
+
+    def __init__(self, datasets, weights):
+
+        self.datasets = datasets
+        num_datasets = len(datasets)
+        assert num_datasets == len(weights)
+
+        self.size = 0
+        for dataset in self.datasets:
+            self.size += len(dataset)
+
+        # Normalize weights.
+        weights = np.array(weights, dtype=np.float64)
+        sum_weights = np.sum(weights)
+        assert sum_weights > 0.0
+        weights /= sum_weights
+
+        # Build indecies.
+        start_time = time.time()
+        assert num_datasets < 255
+        self.dataset_index = np.zeros(self.size, dtype=np.uint8)
+        self.dataset_sample_index = np.zeros(self.size, dtype=np.int64)
+
+        from fengshen.data.megatron_dataloader import helpers
+        helpers.build_blending_indices(self.dataset_index,
+                                       self.dataset_sample_index,
+                                       weights, num_datasets, self.size,
+                                       torch.distributed.get_rank() == 0)
+        print_rank_0('> elapsed time for building blendable dataset indices: '
+                     '{:.2f} (sec)'.format(time.time() - start_time))
+
+    def __len__(self):
+        return self.size
+
+    def __getitem__(self, idx):
+        dataset_idx = self.dataset_index[idx]
+        sample_idx = self.dataset_sample_index[idx]
+        return self.datasets[dataset_idx][sample_idx]

fengshen/data/megatron_dataloader/dataset_utils.py

@@ -0,0 +1,788 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors, and NVIDIA.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+# Most of the code here has been copied from:
+#   https://github.com/google-research/albert/blob/master/create_pretraining_data.py
+# with some modifications.
+
+import math
+import time
+import collections
+
+import numpy as np
+import re
+
+from fengshen.data.megatron_dataloader.utils import (
+    print_rank_0
+)
+from fengshen.data.megatron_dataloader.blendable_dataset import BlendableDataset
+from fengshen.data.megatron_dataloader.indexed_dataset import make_dataset as make_indexed_dataset
+
+DSET_TYPE_BERT = 'standard_bert'
+DSET_TYPE_ICT = 'ict'
+DSET_TYPE_T5 = 't5'
+DSET_TYPE_BERT_CN_WWM = 'bert_cn_wwm'
+DSET_TYPE_BART = 'bart'
+DSET_TYPE_COCOLM = 'coco_lm'
+
+DSET_TYPES = [DSET_TYPE_BERT, DSET_TYPE_ICT,
+              DSET_TYPE_T5, DSET_TYPE_BERT_CN_WWM,
+              DSET_TYPE_BART, DSET_TYPE_COCOLM]
+
+
+def get_datasets_weights_and_num_samples(data_prefix,
+                                         train_valid_test_num_samples):
+
+    # The data prefix should be in the format of:
+    #   weight-1, data-prefix-1, weight-2, data-prefix-2, ..
+    assert len(data_prefix) % 2 == 0
+    num_datasets = len(data_prefix) // 2
+    weights = [0] * num_datasets
+    prefixes = [0] * num_datasets
+    for i in range(num_datasets):
+        weights[i] = float(data_prefix[2 * i])
+        prefixes[i] = (data_prefix[2 * i + 1]).strip()
+    # Normalize weights
+    weight_sum = 0.0
+    for weight in weights:
+        weight_sum += weight
+    assert weight_sum > 0.0
+    weights = [weight / weight_sum for weight in weights]
+
+    # Add 0.5% (the 1.005 factor) so in case the bleding dataset does
+    # not uniformly distribute the number of samples, we still have
+    # samples left to feed to the network.
+    datasets_train_valid_test_num_samples = []
+    for weight in weights:
+        datasets_train_valid_test_num_samples.append(
+            [int(math.ceil(val * weight * 1.005))
+             for val in train_valid_test_num_samples])
+
+    return prefixes, weights, datasets_train_valid_test_num_samples
+
+
+def compile_helper():
+    """Compile helper function ar runtime. Make sure this
+    is invoked on a single process."""
+    import os
+    import subprocess
+    path = os.path.abspath(os.path.dirname(__file__))
+    ret = subprocess.run(['make', '-C', path])
+    if ret.returncode != 0:
+        print("Making C++ dataset helpers module failed, exiting.")
+        import sys
+        sys.exit(1)
+
+
+def get_a_and_b_segments(sample, np_rng):
+    """Divide sample into a and b segments."""
+
+    # Number of sentences in the sample.
+    n_sentences = len(sample)
+    # Make sure we always have two sentences.
+    assert n_sentences > 1, 'make sure each sample has at least two sentences.'
+
+    # First part:
+    # `a_end` is how many sentences go into the `A`.
+    a_end = 1
+    if n_sentences >= 3:
+        # Note that randin in numpy is exclusive.
+        a_end = np_rng.randint(1, n_sentences)
+    tokens_a = []
+    for j in range(a_end):
+        tokens_a.extend(sample[j])
+
+    # Second part:
+    tokens_b = []
+    for j in range(a_end, n_sentences):
+        tokens_b.extend(sample[j])
+
+    # Random next:
+    is_next_random = False
+    if np_rng.random() < 0.5:
+        is_next_random = True
+        tokens_a, tokens_b = tokens_b, tokens_a
+
+    return tokens_a, tokens_b, is_next_random
+
+
+def truncate_segments(tokens_a, tokens_b, len_a, len_b, max_num_tokens, np_rng):
+    """Truncates a pair of sequences to a maximum sequence length."""
+    # print(len_a, len_b, max_num_tokens)
+    assert len_a > 0
+    if len_a + len_b <= max_num_tokens:
+        return False
+    while len_a + len_b > max_num_tokens:
+        if len_a > len_b:
+            len_a -= 1
+            tokens = tokens_a
+        else:
+            len_b -= 1
+            tokens = tokens_b
+        if np_rng.random() < 0.5:
+            del tokens[0]
+        else:
+            tokens.pop()
+    return True
+
+
+def create_tokens_and_tokentypes(tokens_a, tokens_b, cls_id, sep_id):
+    """Merge segments A and B, add [CLS] and [SEP] and build tokentypes."""
+
+    tokens = []
+    tokentypes = []
+    # [CLS].
+    tokens.append(cls_id)
+    tokentypes.append(0)
+    # Segment A.
+    for token in tokens_a:
+        tokens.append(token)
+        tokentypes.append(0)
+    # [SEP].
+    tokens.append(sep_id)
+    tokentypes.append(0)
+    # Segment B.
+    for token in tokens_b:
+        tokens.append(token)
+        tokentypes.append(1)
+    if tokens_b:
+        # [SEP].
+        tokens.append(sep_id)
+        tokentypes.append(1)
+
+    return tokens, tokentypes
+
+
+MaskedLmInstance = collections.namedtuple("MaskedLmInstance",
+                                          ["index", "label"])
+
+
+def is_start_piece(piece):
+    """Check if the current word piece is the starting piece (BERT)."""
+    # When a word has been split into
+    # WordPieces, the first token does not have any marker and any subsequence
+    # tokens are prefixed with ##. So whenever we see the ## token, we
+    # append it to the previous set of word indexes.
+    return not piece.startswith("##")
+
+
+def create_masked_lm_predictions(tokens,
+                                 vocab_id_list, vocab_id_to_token_dict,
+                                 masked_lm_prob,
+                                 cls_id, sep_id, mask_id,
+                                 max_predictions_per_seq,
+                                 np_rng,
+                                 tokenizer,
+                                 max_ngrams=3,
+                                 do_whole_word_mask=True,
+                                 favor_longer_ngram=False,
+                                 do_permutation=False,
+                                 geometric_dist=False,
+                                 masking_style="bert",
+                                 zh_tokenizer=None):
+    """Creates the predictions for the masked LM objective.
+    Note: Tokens here are vocab ids and not text tokens."""
+
+    cand_indexes = []
+    # Note(mingdachen): We create a list for recording if the piece is
+    # the starting piece of current token, where 1 means true, so that
+    # on-the-fly whole word masking is possible.
+    token_boundary = [0] * len(tokens)
+
+    # 如果没有指定中文分词器，那就直接按##算
+    if zh_tokenizer is None:
+        for (i, token) in enumerate(tokens):
+            if token == cls_id or token == sep_id:
+                token_boundary[i] = 1
+                continue
+        # Whole Word Masking means that if we mask all of the wordpieces
+        # corresponding to an original word.
+        #
+        # Note that Whole Word Masking does *not* change the training code
+        # at all -- we still predict each WordPiece independently, softmaxed
+        # over the entire vocabulary.
+            if (do_whole_word_mask and len(cand_indexes) >= 1 and
+                    not is_start_piece(vocab_id_to_token_dict[token])):
+                cand_indexes[-1].append(i)
+            else:
+                cand_indexes.append([i])
+                if is_start_piece(vocab_id_to_token_dict[token]):
+                    token_boundary[i] = 1
+    else:
+        # 如果指定了中文分词器，那就先用分词器分词，然后再进行判断
+        # 获取去掉CLS SEP的原始文本
+        raw_tokens = []
+        for t in tokens:
+            if t != cls_id and t != sep_id:
+                raw_tokens.append(t)
+        raw_tokens = [vocab_id_to_token_dict[i] for i in raw_tokens]
+        # 分词然后获取每次字开头的最长词的长度
+        word_list = set(zh_tokenizer(''.join(raw_tokens), HMM=True))
+        word_length_dict = {}
+        for w in word_list:
+            if len(w) < 1:
+                continue
+            if w[0] not in word_length_dict:
+                word_length_dict[w[0]] = len(w)
+            elif word_length_dict[w[0]] < len(w):
+                word_length_dict[w[0]] = len(w)
+        i = 0
+        # 从词表里面检索
+        while i < len(tokens):
+            token_id = tokens[i]
+            token = vocab_id_to_token_dict[token_id]
+            if len(token) == 0 or token_id == cls_id or token_id == sep_id:
+                token_boundary[i] = 1
+                i += 1
+                continue
+            word_max_length = 1
+            if token[0] in word_length_dict:
+                word_max_length = word_length_dict[token[0]]
+            j = 0
+            word = ''
+            word_end = i+1
+            # 兼容以前##的形式，如果后面的词是##开头的，那么直接把后面的拼到前面当作一个词
+            old_style = False
+            while word_end < len(tokens) and vocab_id_to_token_dict[tokens[word_end]].startswith('##'):
+                old_style = True
+                word_end += 1
+            if not old_style:
+                while j < word_max_length and i+j < len(tokens):
+                    cur_token = tokens[i+j]
+                    word += vocab_id_to_token_dict[cur_token]
+                    j += 1
+                    if word in word_list:
+                        word_end = i+j
+            cand_indexes.append([p for p in range(i, word_end)])
+            token_boundary[i] = 1
+            i = word_end
+
+    output_tokens = list(tokens)
+    # add by ganruyi
+    if masking_style == 'bert-cn-wwm':
+        # if non chinese is False, that means it is chinese
+        # then try to remove "##" which is added previously
+        new_token_ids = []
+        for token_id in output_tokens:
+            token = tokenizer.convert_ids_to_tokens([token_id])[0]
+            if len(re.findall('##[\u4E00-\u9FA5]', token)) > 0:
+                token = token[2:]
+            new_token_id = tokenizer.convert_tokens_to_ids([token])[
+                0]
+            new_token_ids.append(new_token_id)
+        output_tokens = new_token_ids
+
+    masked_lm_positions = []
+    masked_lm_labels = []
+
+    if masked_lm_prob == 0:
+        return (output_tokens, masked_lm_positions,
+                masked_lm_labels, token_boundary)
+
+    num_to_predict = min(max_predictions_per_seq,
+                         max(1, int(round(len(tokens) * masked_lm_prob))))
+
+    ngrams = np.arange(1, max_ngrams + 1, dtype=np.int64)
+    if not geometric_dist:
+        # Note(mingdachen):
+        # By default, we set the probilities to favor shorter ngram sequences.
+        pvals = 1. / np.arange(1, max_ngrams + 1)
+        pvals /= pvals.sum(keepdims=True)
+        if favor_longer_ngram:
+            pvals = pvals[::-1]
+    # 获取一个ngram的idx，对于每个word，记录他的ngram的word
+    ngram_indexes = []
+    for idx in range(len(cand_indexes)):
+        ngram_index = []
+        for n in ngrams:
+            ngram_index.append(cand_indexes[idx:idx + n])
+        ngram_indexes.append(ngram_index)
+
+    np_rng.shuffle(ngram_indexes)
+
+    (masked_lms, masked_spans) = ([], [])
+    covered_indexes = set()
+    for cand_index_set in ngram_indexes:
+        if len(masked_lms) >= num_to_predict:
+            break
+        if not cand_index_set:
+            continue
+        # Note(mingdachen):
+        # Skip current piece if they are covered in lm masking or previous ngrams.
+        for index_set in cand_index_set[0]:
+            for index in index_set:
+                if index in covered_indexes:
+                    continue
+
+        if not geometric_dist:
+            n = np_rng.choice(ngrams[:len(cand_index_set)],
+                              p=pvals[:len(cand_index_set)] /
+                              pvals[:len(cand_index_set)].sum(keepdims=True))
+        else:
+            # Sampling "n" from the geometric distribution and clipping it to
+            # the max_ngrams. Using p=0.2 default from the SpanBERT paper
+            # https://arxiv.org/pdf/1907.10529.pdf (Sec 3.1)
+            n = min(np_rng.geometric(0.2), max_ngrams)
+
+        index_set = sum(cand_index_set[n - 1], [])
+        n -= 1
+        # Note(mingdachen):
+        # Repeatedly looking for a candidate that does not exceed the
+        # maximum number of predictions by trying shorter ngrams.
+        while len(masked_lms) + len(index_set) > num_to_predict:
+            if n == 0:
+                break
+            index_set = sum(cand_index_set[n - 1], [])
+            n -= 1
+        # If adding a whole-word mask would exceed the maximum number of
+        # predictions, then just skip this candidate.
+        if len(masked_lms) + len(index_set) > num_to_predict:
+            continue
+        is_any_index_covered = False
+        for index in index_set:
+            if index in covered_indexes:
+                is_any_index_covered = True
+                break
+        if is_any_index_covered:
+            continue
+        for index in index_set:
+            covered_indexes.add(index)
+            masked_token = None
+            if masking_style == "bert":
+                # 80% of the time, replace with [MASK]
+                if np_rng.random() < 0.8:
+                    masked_token = mask_id
+                else:
+                    # 10% of the time, keep original
+                    if np_rng.random() < 0.5:
+                        masked_token = tokens[index]
+                    # 10% of the time, replace with random word
+                    else:
+                        masked_token = vocab_id_list[np_rng.randint(0, len(vocab_id_list))]
+            elif masking_style == 'bert-cn-wwm':
+                # 80% of the time, replace with [MASK]
+                if np_rng.random() < 0.8:
+                    masked_token = mask_id
+                else:
+                    # 10% of the time, keep original
+                    if np_rng.random() < 0.5:
+                        # 如果是中文全词mask，去掉tokens里的##
+                        token_id = tokens[index]
+                        token = tokenizer.convert_ids_to_tokens([token_id])[
+                            0]
+                        if len(re.findall('##[\u4E00-\u9FA5]', token)) > 0:
+                            token = token[2:]
+                        new_token_id = tokenizer.convert_tokens_to_ids([token])[
+                            0]
+                        masked_token = new_token_id
+                    # 10% of the time, replace with random word
+                    else:
+                        masked_token = vocab_id_list[np_rng.randint(
+                            0, len(vocab_id_list))]
+            elif masking_style == "t5":
+                masked_token = mask_id
+            else:
+                raise ValueError("invalid value of masking style")
+
+            output_tokens[index] = masked_token
+            masked_lms.append(MaskedLmInstance(
+                index=index, label=tokens[index]))
+
+        masked_spans.append(MaskedLmInstance(
+            index=index_set,
+            label=[tokens[index] for index in index_set]))
+
+    assert len(masked_lms) <= num_to_predict
+    np_rng.shuffle(ngram_indexes)
+
+    select_indexes = set()
+    if do_permutation:
+        for cand_index_set in ngram_indexes:
+            if len(select_indexes) >= num_to_predict:
+                break
+            if not cand_index_set:
+                continue
+            # Note(mingdachen):
+            # Skip current piece if they are covered in lm masking or previous ngrams.
+            for index_set in cand_index_set[0]:
+                for index in index_set:
+                    if index in covered_indexes or index in select_indexes:
+                        continue
+
+            n = np.random.choice(ngrams[:len(cand_index_set)],
+                                 p=pvals[:len(cand_index_set)] /
+                                 pvals[:len(cand_index_set)].sum(keepdims=True))
+            index_set = sum(cand_index_set[n - 1], [])
+            n -= 1
+
+            while len(select_indexes) + len(index_set) > num_to_predict:
+                if n == 0:
+                    break
+                index_set = sum(cand_index_set[n - 1], [])
+                n -= 1
+            # If adding a whole-word mask would exceed the maximum number of
+            # predictions, then just skip this candidate.
+            if len(select_indexes) + len(index_set) > num_to_predict:
+                continue
+            is_any_index_covered = False
+            for index in index_set:
+                if index in covered_indexes or index in select_indexes:
+                    is_any_index_covered = True
+                    break
+            if is_any_index_covered:
+                continue
+            for index in index_set:
+                select_indexes.add(index)
+        assert len(select_indexes) <= num_to_predict
+
+        select_indexes = sorted(select_indexes)
+        permute_indexes = list(select_indexes)
+        np_rng.shuffle(permute_indexes)
+        orig_token = list(output_tokens)
+
+        for src_i, tgt_i in zip(select_indexes, permute_indexes):
+            output_tokens[src_i] = orig_token[tgt_i]
+            masked_lms.append(MaskedLmInstance(
+                index=src_i, label=orig_token[src_i]))
+
+    masked_lms = sorted(masked_lms, key=lambda x: x.index)
+    # Sort the spans by the index of the first span
+    masked_spans = sorted(masked_spans, key=lambda x: x.index[0])
+
+    for p in masked_lms:
+        masked_lm_positions.append(p.index)
+        masked_lm_labels.append(p.label)
+    return (output_tokens, masked_lm_positions, masked_lm_labels, token_boundary, masked_spans)
+
+
+def pad_and_convert_to_numpy(tokens, tokentypes, masked_positions,
+                             masked_labels, pad_id, max_seq_length):
+    """Pad sequences and convert them to numpy."""
+
+    # Some checks.
+    num_tokens = len(tokens)
+    padding_length = max_seq_length - num_tokens
+    assert padding_length >= 0
+    assert len(tokentypes) == num_tokens
+    assert len(masked_positions) == len(masked_labels)
+
+    # Tokens and token types.
+    filler = [pad_id] * padding_length
+    tokens_np = np.array(tokens + filler, dtype=np.int64)
+    tokentypes_np = np.array(tokentypes + filler, dtype=np.int64)
+
+    # Padding mask.
+    padding_mask_np = np.array([1] * num_tokens + [0] * padding_length,
+                               dtype=np.int64)
+
+    # Lables and loss mask.
+    labels = [-1] * max_seq_length
+    loss_mask = [0] * max_seq_length
+    for i in range(len(masked_positions)):
+        assert masked_positions[i] < num_tokens
+        labels[masked_positions[i]] = masked_labels[i]
+        loss_mask[masked_positions[i]] = 1
+    labels_np = np.array(labels, dtype=np.int64)
+    loss_mask_np = np.array(loss_mask, dtype=np.int64)
+
+    return tokens_np, tokentypes_np, labels_np, padding_mask_np, loss_mask_np
+
+
+def build_train_valid_test_datasets(data_prefix, data_impl, splits_string,
+                                    train_valid_test_num_samples,
+                                    max_seq_length,
+                                    masked_lm_prob, short_seq_prob, seed,
+                                    tokenizer,
+                                    skip_warmup, binary_head=False,
+                                    max_seq_length_dec=None,
+                                    dataset_type='standard_bert',
+                                    zh_tokenizer=None,
+                                    span=None):
+
+    if len(data_prefix) == 1:
+        return _build_train_valid_test_datasets(data_prefix[0],
+                                                data_impl, splits_string,
+                                                train_valid_test_num_samples,
+                                                max_seq_length, masked_lm_prob,
+                                                short_seq_prob, seed,
+                                                skip_warmup,
+                                                binary_head,
+                                                max_seq_length_dec,
+                                                tokenizer,
+                                                dataset_type=dataset_type,
+                                                zh_tokenizer=zh_tokenizer,
+                                                span=span)
+    # Blending dataset.
+    # Parse the values.
+    output = get_datasets_weights_and_num_samples(data_prefix,
+                                                  train_valid_test_num_samples)
+    prefixes, weights, datasets_train_valid_test_num_samples = output
+
+    # Build individual datasets.
+    train_datasets = []
+    valid_datasets = []
+    test_datasets = []
+    for i in range(len(prefixes)):
+        train_ds, valid_ds, test_ds = _build_train_valid_test_datasets(
+            prefixes[i], data_impl, splits_string,
+            datasets_train_valid_test_num_samples[i],
+            max_seq_length, masked_lm_prob, short_seq_prob,
+            seed, skip_warmup, binary_head, max_seq_length_dec,
+            tokenizer, dataset_type=dataset_type, zh_tokenizer=zh_tokenizer)
+        if train_ds:
+            train_datasets.append(train_ds)
+        if valid_ds:
+            valid_datasets.append(valid_ds)
+        if test_ds:
+            test_datasets.append(test_ds)
+
+        # Blend.
+    blending_train_dataset = None
+    if train_datasets:
+        blending_train_dataset = BlendableDataset(train_datasets, weights)
+    blending_valid_dataset = None
+    if valid_datasets:
+        blending_valid_dataset = BlendableDataset(valid_datasets, weights)
+    blending_test_dataset = None
+    if test_datasets:
+        blending_test_dataset = BlendableDataset(test_datasets, weights)
+
+    return (blending_train_dataset, blending_valid_dataset,
+            blending_test_dataset)
+
+
+def _build_train_valid_test_datasets(data_prefix, data_impl, splits_string,
+                                     train_valid_test_num_samples,
+                                     max_seq_length,
+                                     masked_lm_prob, short_seq_prob, seed,
+                                     skip_warmup, binary_head,
+                                     max_seq_length_dec,
+                                     tokenizer,
+                                     dataset_type='standard_bert',
+                                     zh_tokenizer=None,
+                                     span=None):
+
+    if dataset_type not in DSET_TYPES:
+        raise ValueError("Invalid dataset_type: ", dataset_type)
+
+    # Indexed dataset.
+    indexed_dataset = get_indexed_dataset_(data_prefix,
+                                           data_impl,
+                                           skip_warmup)
+
+    # Get start and end indices of train/valid/train into doc-idx
+    # Note that doc-idx is desinged to be num-docs + 1 so we can
+    # easily iterate over it.
+    total_num_of_documents = indexed_dataset.doc_idx.shape[0] - 1
+    splits = get_train_valid_test_split_(splits_string, total_num_of_documents)
+
+    # Print stats about the splits.
+    print_rank_0(' > dataset split:')
+
+    def print_split_stats(name, index):
+        print_rank_0('    {}:'.format(name))
+        print_rank_0('     document indices in [{}, {}) total of {} '
+                     'documents'.format(splits[index], splits[index + 1],
+                                        splits[index + 1] - splits[index]))
+        start_index = indexed_dataset.doc_idx[splits[index]]
+        end_index = indexed_dataset.doc_idx[splits[index + 1]]
+        print_rank_0('     sentence indices in [{}, {}) total of {} '
+                     'sentences'.format(start_index, end_index,
+                                        end_index - start_index))
+    print_split_stats('train', 0)
+    print_split_stats('validation', 1)
+    print_split_stats('test', 2)
+
+    def build_dataset(index, name):
+        from fengshen.data.megatron_dataloader.bert_dataset import BertDataset
+        from fengshen.data.megatron_dataloader.bart_dataset import BartDataset
+        from fengshen.data.megatron_dataloader.cocolm_dataset import COCOLMDataset
+        dataset = None
+        if splits[index + 1] > splits[index]:
+            # Get the pointer to the original doc-idx so we can set it later.
+            doc_idx_ptr = indexed_dataset.get_doc_idx()
+            # Slice the doc-idx
+            start_index = splits[index]
+            # Add +1 so we can index into the dataset to get the upper bound.
+            end_index = splits[index + 1] + 1
+            # New doc_idx view.
+            indexed_dataset.set_doc_idx(doc_idx_ptr[start_index:end_index])
+            # Build the dataset accordingly.
+            kwargs = dict(
+                name=name,
+                data_prefix=data_prefix,
+                num_epochs=None,
+                max_num_samples=train_valid_test_num_samples[index],
+                max_seq_length=max_seq_length,
+                seed=seed,
+            )
+
+            if dataset_type == DSET_TYPE_BERT or dataset_type == DSET_TYPE_BERT_CN_WWM:
+                dataset = BertDataset(
+                    indexed_dataset=indexed_dataset,
+                    masked_lm_prob=masked_lm_prob,
+                    short_seq_prob=short_seq_prob,
+                    binary_head=binary_head,
+                    # 增加参数区分bert和bert-cn-wwm
+                    tokenizer=tokenizer,
+                    masking_style='bert' if dataset_type == DSET_TYPE_BERT else 'bert-cn-wwm',
+                    **kwargs
+                )
+            elif dataset_type == DSET_TYPE_BART:
+                dataset = BartDataset(
+                    indexed_dataset=indexed_dataset,
+                    masked_lm_prob=masked_lm_prob,
+                    short_seq_prob=short_seq_prob,
+                    tokenizer=tokenizer,
+                    zh_tokenizer=zh_tokenizer,
+                    **kwargs
+                )
+            elif dataset_type == DSET_TYPE_COCOLM:
+                dataset = COCOLMDataset(
+                    indexed_dataset=indexed_dataset,
+                    masked_lm_prob=masked_lm_prob,
+                    short_seq_prob=short_seq_prob,
+                    tokenizer=tokenizer,
+                    masking_style='bert',
+                    span=span,
+                    **kwargs
+                )
+            else:
+                raise NotImplementedError(
+                    "Dataset type not fully implemented.")
+
+            # Set the original pointer so dataset remains the main dataset.
+            indexed_dataset.set_doc_idx(doc_idx_ptr)
+            # Checks.
+            assert indexed_dataset.doc_idx[0] == 0
+            assert indexed_dataset.doc_idx.shape[0] == \
+                (total_num_of_documents + 1)
+        return dataset
+
+    train_dataset = build_dataset(0, 'train')
+    valid_dataset = build_dataset(1, 'valid')
+    test_dataset = build_dataset(2, 'test')
+
+    return (train_dataset, valid_dataset, test_dataset)
+
+
+def get_indexed_dataset_(data_prefix, data_impl, skip_warmup):
+
+    print_rank_0(' > building dataset index ...')
+
+    start_time = time.time()
+    indexed_dataset = make_indexed_dataset(data_prefix,
+                                           data_impl,
+                                           skip_warmup)
+    assert indexed_dataset.sizes.shape[0] == indexed_dataset.doc_idx[-1]
+    print_rank_0(' > finished creating indexed dataset in {:4f} '
+                 'seconds'.format(time.time() - start_time))
+
+    print_rank_0(' > indexed dataset stats:')
+    print_rank_0('    number of documents: {}'.format(
+        indexed_dataset.doc_idx.shape[0] - 1))
+    print_rank_0('    number of sentences: {}'.format(
+        indexed_dataset.sizes.shape[0]))
+
+    return indexed_dataset
+
+
+def get_train_valid_test_split_(splits_string, size):
+    """ Get dataset splits from comma or '/' separated string list."""
+
+    splits = []
+    if splits_string.find(',') != -1:
+        splits = [float(s) for s in splits_string.split(',')]
+    elif splits_string.find('/') != -1:
+        splits = [float(s) for s in splits_string.split('/')]
+    else:
+        splits = [float(splits_string)]
+    while len(splits) < 3:
+        splits.append(0.)
+    splits = splits[:3]
+    splits_sum = sum(splits)
+    assert splits_sum > 0.0
+    splits = [split / splits_sum for split in splits]
+    splits_index = [0]
+    for index, split in enumerate(splits):
+        splits_index.append(splits_index[index] +
+                            int(round(split * float(size))))
+    diff = splits_index[-1] - size
+    for index in range(1, len(splits_index)):
+        splits_index[index] -= diff
+    assert len(splits_index) == 4
+    assert splits_index[-1] == size
+    return splits_index
+
+
+def get_samples_mapping(indexed_dataset,
+                        data_prefix,
+                        num_epochs,
+                        max_num_samples,
+                        max_seq_length,
+                        short_seq_prob,
+                        seed,
+                        name,
+                        binary_head):
+    """Get a list that maps a sample index to a starting
+    sentence index, end sentence index, and length"""
+
+    if not num_epochs:
+        if not max_num_samples:
+            raise ValueError("Need to specify either max_num_samples "
+                             "or num_epochs")
+        num_epochs = np.iinfo(np.int32).max - 1
+    if not max_num_samples:
+        max_num_samples = np.iinfo(np.int64).max - 1
+
+    # Filename of the index mapping
+    indexmap_filename = data_prefix
+    indexmap_filename += '_{}_indexmap'.format(name)
+    if num_epochs != (np.iinfo(np.int32).max - 1):
+        indexmap_filename += '_{}ep'.format(num_epochs)
+    if max_num_samples != (np.iinfo(np.int64).max - 1):
+        indexmap_filename += '_{}mns'.format(max_num_samples)
+    indexmap_filename += '_{}msl'.format(max_seq_length)
+    indexmap_filename += '_{:0.2f}ssp'.format(short_seq_prob)
+    indexmap_filename += '_{}s'.format(seed)
+    indexmap_filename += '.npy'
+
+    # This should be a barrier but nccl barrier assumes
+    # device_index=rank which is not the case for model
+    # parallel case
+    # ganruyi comment
+    # counts = torch.cuda.LongTensor([1])
+    # torch.distributed.all_reduce(
+    # counts, group=mpu.get_data_parallel_group())
+    # torch.distributed.all_reduce(
+    # counts, group=mpu.get_pipeline_model_parallel_group())
+    # assert counts[0].item() == (
+    #    torch.distributed.get_world_size() //
+    #    torch.distributed.get_world_size(
+    # group=mpu.get_tensor_model_parallel_group()))
+
+    # Load indexed dataset.
+    print_rank_0(' > loading indexed mapping from {}'.format(
+        indexmap_filename))
+    start_time = time.time()
+    samples_mapping = np.load(
+        indexmap_filename, allow_pickle=True, mmap_mode='r')
+    print_rank_0('    loaded indexed file in {:3.3f} seconds'.format(
+        time.time() - start_time))
+    print_rank_0('    total number of samples: {}'.format(
+        samples_mapping.shape[0]))
+
+    return samples_mapping

fengshen/data/megatron_dataloader/helpers.cpp

@@ -0,0 +1,794 @@
+/*
+ coding=utf-8
+ Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+ */
+
+/* Helper methods for fast index mapping builds */
+
+#include <algorithm>
+#include <iostream>
+#include <limits>
+#include <math.h>
+#include <stdexcept>
+#include <pybind11/pybind11.h>
+#include <pybind11/numpy.h>
+#include <random>
+
+namespace py = pybind11;
+using namespace std;
+
+const int32_t LONG_SENTENCE_LEN = 512;
+
+void build_blending_indices(py::array_t<uint8_t> &dataset_index,
+                            py::array_t<int64_t> &dataset_sample_index,
+                            const py::array_t<double> &weights,
+                            const int32_t num_datasets,
+                            const int64_t size, const bool verbose)
+{
+  /* Given multiple datasets and a weighting array, build samples
+   such that it follows those wieghts.*/
+
+  if (verbose)
+  {
+    std::cout << "> building indices for blendable datasets ..." << std::endl;
+  }
+
+  // Get the pointer access without the checks.
+  auto dataset_index_ptr = dataset_index.mutable_unchecked<1>();
+  auto dataset_sample_index_ptr = dataset_sample_index.mutable_unchecked<1>();
+  auto weights_ptr = weights.unchecked<1>();
+
+  // Initialize buffer for number of samples used for each dataset.
+  int64_t current_samples[num_datasets];
+  for (int64_t i = 0; i < num_datasets; ++i)
+  {
+    current_samples[i] = 0;
+  }
+
+  // For each sample:
+  for (int64_t sample_idx = 0; sample_idx < size; ++sample_idx)
+  {
+
+    // Determine where the max error in sampling is happening.
+    auto sample_idx_double = std::max(static_cast<double>(sample_idx), 1.0);
+    int64_t max_error_index = 0;
+    double max_error = weights_ptr[0] * sample_idx_double -
+                       static_cast<double>(current_samples[0]);
+    for (int64_t dataset_idx = 1; dataset_idx < num_datasets; ++dataset_idx)
+    {
+      double error = weights_ptr[dataset_idx] * sample_idx_double -
+                     static_cast<double>(current_samples[dataset_idx]);
+      if (error > max_error)
+      {
+        max_error = error;
+        max_error_index = dataset_idx;
+      }
+    }
+
+    // Populate the indices.
+    dataset_index_ptr[sample_idx] = static_cast<uint8_t>(max_error_index);
+    dataset_sample_index_ptr[sample_idx] = current_samples[max_error_index];
+
+    // Update the total samples.
+    current_samples[max_error_index] += 1;
+  }
+
+  // print info
+  if (verbose)
+  {
+    std::cout << " > sample ratios:" << std::endl;
+    for (int64_t dataset_idx = 0; dataset_idx < num_datasets; ++dataset_idx)
+    {
+      auto ratio = static_cast<double>(current_samples[dataset_idx]) /
+                   static_cast<double>(size);
+      std::cout << "   dataset " << dataset_idx << ", input: " << weights_ptr[dataset_idx] << ", achieved: " << ratio << std::endl;
+    }
+  }
+}
+
+py::array build_sample_idx(const py::array_t<int32_t> &sizes_,
+                           const py::array_t<int32_t> &doc_idx_,
+                           const int32_t seq_length,
+                           const int32_t num_epochs,
+                           const int64_t tokens_per_epoch)
+{
+  /* Sample index (sample_idx) is used for gpt2 like dataset for which
+       the documents are flattened and the samples are built based on this
+       1-D flatten array. It is a 2D array with sizes [number-of-samples + 1, 2]
+       where [..., 0] contains the index into `doc_idx` and [..., 1] is the
+       starting offset in that document.*/
+
+  // Consistency checks.
+  assert(seq_length > 1);
+  assert(num_epochs > 0);
+  assert(tokens_per_epoch > 1);
+
+  // Remove bound checks.
+  auto sizes = sizes_.unchecked<1>();
+  auto doc_idx = doc_idx_.unchecked<1>();
+
+  // Mapping and it's length (1D).
+  int64_t num_samples = (num_epochs * tokens_per_epoch - 1) / seq_length;
+  int32_t *sample_idx = new int32_t[2 * (num_samples + 1)];
+
+  cout << "    using:" << endl
+       << std::flush;
+  cout << "     number of documents:       " << doc_idx_.shape(0) / num_epochs << endl
+       << std::flush;
+  cout << "     number of epochs:          " << num_epochs << endl
+       << std::flush;
+  cout << "     sequence length:           " << seq_length << endl
+       << std::flush;
+  cout << "     total number of samples:   " << num_samples << endl
+       << std::flush;
+
+  // Index into sample_idx.
+  int64_t sample_index = 0;
+  // Index into doc_idx.
+  int64_t doc_idx_index = 0;
+  // Begining offset for each document.
+  int32_t doc_offset = 0;
+  // Start with first document and no offset.
+  sample_idx[2 * sample_index] = doc_idx_index;
+  sample_idx[2 * sample_index + 1] = doc_offset;
+  ++sample_index;
+
+  while (sample_index <= num_samples)
+  {
+    // Start with a fresh sequence.
+    int32_t remaining_seq_length = seq_length + 1;
+    while (remaining_seq_length != 0)
+    {
+      // Get the document length.
+      auto doc_id = doc_idx[doc_idx_index];
+      auto doc_length = sizes[doc_id] - doc_offset;
+      // And add it to the current sequence.
+      remaining_seq_length -= doc_length;
+      // If we have more than a full sequence, adjust offset and set
+      // remaining length to zero so we return from the while loop.
+      // Note that -1 here is for the same reason we have -1 in
+      // `_num_epochs` calculations.
+      if (remaining_seq_length <= 0)
+      {
+        doc_offset += (remaining_seq_length + doc_length - 1);
+        remaining_seq_length = 0;
+      }
+      else
+      {
+        // Otherwise, start from the begining of the next document.
+        ++doc_idx_index;
+        doc_offset = 0;
+      }
+    }
+    // Record the sequence.
+    sample_idx[2 * sample_index] = doc_idx_index;
+    sample_idx[2 * sample_index + 1] = doc_offset;
+    ++sample_index;
+  }
+
+  // Method to deallocate memory.
+  py::capsule free_when_done(sample_idx, [](void *mem_)
+                             {
+                               int32_t *mem = reinterpret_cast<int32_t *>(mem_);
+                               delete[] mem;
+                             });
+
+  // Return the numpy array.
+  const auto byte_size = sizeof(int32_t);
+  return py::array(std::vector<int64_t>{num_samples + 1, 2}, // shape
+                   {2 * byte_size, byte_size},               // C-style contiguous strides
+                   sample_idx,                               // the data pointer
+                   free_when_done);                          // numpy array references
+}
+
+inline int32_t get_target_sample_len(const int32_t short_seq_ratio,
+                                     const int32_t max_length,
+                                     std::mt19937 &rand32_gen)
+{
+  /* Training sample length. */
+  if (short_seq_ratio == 0)
+  {
+    return max_length;
+  }
+  const auto random_number = rand32_gen();
+  if ((random_number % short_seq_ratio) == 0)
+  {
+    return 2 + random_number % (max_length - 1);
+  }
+  return max_length;
+}
+
+template <typename DocIdx>
+py::array build_mapping_impl(const py::array_t<int64_t> &docs_,
+                             const py::array_t<int32_t> &sizes_,
+                             const int32_t num_epochs,
+                             const uint64_t max_num_samples,
+                             const int32_t max_seq_length,
+                             const double short_seq_prob,
+                             const int32_t seed,
+                             const bool verbose,
+                             const int32_t min_num_sent)
+{
+  /* Build a mapping of (start-index, end-index, sequence-length) where
+       start and end index are the indices of the sentences in the sample
+       and sequence-length is the target sequence length.
+    */
+
+  // Consistency checks.
+  assert(num_epochs > 0);
+  assert(max_seq_length > 1);
+  assert(short_seq_prob >= 0.0);
+  assert(short_seq_prob <= 1.0);
+  assert(seed > 0);
+
+  // Remove bound checks.
+  auto docs = docs_.unchecked<1>();
+  auto sizes = sizes_.unchecked<1>();
+
+  // For efficiency, convert probability to ratio. Note: rand() generates int.
+  int32_t short_seq_ratio = 0;
+  if (short_seq_prob > 0)
+  {
+    short_seq_ratio = static_cast<int32_t>(round(1.0 / short_seq_prob));
+  }
+
+  if (verbose)
+  {
+    const auto sent_start_index = docs[0];
+    const auto sent_end_index = docs[docs_.shape(0) - 1];
+    const auto num_sentences = sent_end_index - sent_start_index;
+    cout << "    using:" << endl
+         << std::flush;
+    cout << "     number of documents:            " << docs_.shape(0) - 1 << endl
+         << std::flush;
+    cout << "     sentences range:                [" << sent_start_index << ", " << sent_end_index << ")" << endl
+         << std::flush;
+    cout << "     total number of sentences:      " << num_sentences << endl
+         << std::flush;
+    cout << "     number of epochs:               " << num_epochs << endl
+         << std::flush;
+    cout << "     maximum number of samples:      " << max_num_samples << endl
+         << std::flush;
+    cout << "     maximum sequence length:        " << max_seq_length << endl
+         << std::flush;
+    cout << "     short sequence probability:     " << short_seq_prob << endl
+         << std::flush;
+    cout << "     short sequence ration (1/prob): " << short_seq_ratio << endl
+         << std::flush;
+    cout << "     seed:                           " << seed << endl
+         << std::flush;
+  }
+
+  // Mapping and it's length (1D).
+  int64_t num_samples = -1;
+  DocIdx *maps = NULL;
+
+  // Perform two iterations, in the first iteration get the size
+  // and allocate memory and in the second iteration populate the map.
+  bool second = false;
+  for (int32_t iteration = 0; iteration < 2; ++iteration)
+  {
+
+    // Set the seed so both iterations produce the same results.
+    std::mt19937 rand32_gen(seed);
+
+    // Set the flag on second iteration.
+    second = (iteration == 1);
+
+    // Counters:
+    uint64_t empty_docs = 0;
+    uint64_t one_sent_docs = 0;
+    uint64_t long_sent_docs = 0;
+
+    // Current map index.
+    uint64_t map_index = 0;
+
+    // For each epoch:
+    for (int32_t epoch = 0; epoch < num_epochs; ++epoch)
+    {
+      if (map_index >= max_num_samples)
+      {
+        if (verbose && (!second))
+        {
+          cout << "    reached " << max_num_samples << " samples after "
+               << epoch << " epochs ..." << endl
+               << std::flush;
+        }
+        break;
+      }
+      // For each document:
+      for (int32_t doc = 0; doc < (docs.shape(0) - 1); ++doc)
+      {
+
+        // Document sentences are in [sent_index_first, sent_index_last)
+        const auto sent_index_first = docs[doc];
+        const auto sent_index_last = docs[doc + 1];
+
+        // At the begining of the document previous index is the
+        // start index.
+        auto prev_start_index = sent_index_first;
+
+        // Remaining documents.
+        auto num_remain_sent = sent_index_last - sent_index_first;
+
+        // Some bookkeeping
+        if ((epoch == 0) && (!second))
+        {
+          if (num_remain_sent == 0)
+          {
+            ++empty_docs;
+          }
+          if (num_remain_sent == 1)
+          {
+            ++one_sent_docs;
+          }
+        }
+
+        // Detect documents with long sentences.
+        bool contains_long_sentence = false;
+        if (num_remain_sent > 1)
+        {
+          for (auto sent_index = sent_index_first;
+               sent_index < sent_index_last; ++sent_index)
+          {
+            if (sizes[sent_index] > LONG_SENTENCE_LEN)
+            {
+              if ((epoch == 0) && (!second))
+              {
+                ++long_sent_docs;
+              }
+              contains_long_sentence = true;
+              break;
+            }
+          }
+        }
+
+        // If we have more than two sentences.
+        if ((num_remain_sent >= min_num_sent) && (!contains_long_sentence))
+        {
+
+          // Set values.
+          auto seq_len = int32_t{0};
+          auto num_sent = int32_t{0};
+          auto target_seq_len = get_target_sample_len(short_seq_ratio,
+                                                      max_seq_length,
+                                                      rand32_gen);
+
+          // Loop through sentences.
+          for (auto sent_index = sent_index_first;
+               sent_index < sent_index_last; ++sent_index)
+          {
+
+            // Add the size and number of sentences.
+            seq_len += sizes[sent_index];
+            ++num_sent;
+            --num_remain_sent;
+
+            // If we have reached the target length.
+            // and if not only one sentence is left in the document.
+            // and if we have at least two sentneces.
+            // and if we have reached end of the document.
+            if (((seq_len >= target_seq_len) &&
+                 (num_remain_sent > 1) &&
+                 (num_sent >= min_num_sent)) ||
+                (num_remain_sent == 0))
+            {
+
+              // Check for overflow.
+              if ((3 * map_index + 2) >
+                  std::numeric_limits<int64_t>::max())
+              {
+                cout << "number of samples exceeded maximum "
+                     << "allowed by type int64: "
+                     << std::numeric_limits<int64_t>::max()
+                     << endl;
+                throw std::overflow_error("Number of samples");
+              }
+
+              // Populate the map.
+              if (second)
+              {
+                const auto map_index_0 = 3 * map_index;
+                maps[map_index_0] = static_cast<DocIdx>(prev_start_index);
+                maps[map_index_0 + 1] = static_cast<DocIdx>(sent_index + 1);
+                maps[map_index_0 + 2] = static_cast<DocIdx>(target_seq_len);
+              }
+
+              // Update indices / counters.
+              ++map_index;
+              prev_start_index = sent_index + 1;
+              target_seq_len = get_target_sample_len(short_seq_ratio,
+                                                     max_seq_length,
+                                                     rand32_gen);
+              seq_len = 0;
+              num_sent = 0;
+            }
+
+          } // for (auto sent_index=sent_index_first; ...
+        }   // if (num_remain_sent > 1) {
+      }     // for (int doc=0; doc < num_docs; ++doc) {
+    }       // for (int epoch=0; epoch < num_epochs; ++epoch) {
+
+    if (!second)
+    {
+      if (verbose)
+      {
+        cout << "   number of empty documents: " << empty_docs << endl
+             << std::flush;
+        cout << "   number of documents with one sentence: " << one_sent_docs << endl
+             << std::flush;
+        cout << "   number of documents with long sentences: " << long_sent_docs << endl
+             << std::flush;
+        cout << "   will create mapping for " << map_index << " samples" << endl
+             << std::flush;
+      }
+      assert(maps == NULL);
+      assert(num_samples < 0);
+      maps = new DocIdx[3 * map_index];
+      num_samples = static_cast<int64_t>(map_index);
+    }
+
+  } // for (int iteration=0; iteration < 2; ++iteration) {
+
+  // Shuffle.
+  // We need a 64 bit random number generator as we might have more
+  // than 2 billion samples.
+  std::mt19937_64 rand64_gen(seed + 1);
+  for (auto i = (num_samples - 1); i > 0; --i)
+  {
+    const auto j = static_cast<int64_t>(rand64_gen() % (i + 1));
+    const auto i0 = 3 * i;
+    const auto j0 = 3 * j;
+    // Swap values.
+    swap(maps[i0], maps[j0]);
+    swap(maps[i0 + 1], maps[j0 + 1]);
+    swap(maps[i0 + 2], maps[j0 + 2]);
+  }
+
+  // Method to deallocate memory.
+  py::capsule free_when_done(maps, [](void *mem_)
+                             {
+                               DocIdx *mem = reinterpret_cast<DocIdx *>(mem_);
+                               delete[] mem;
+                             });
+
+  // Return the numpy array.
+  const auto byte_size = sizeof(DocIdx);
+  return py::array(std::vector<int64_t>{num_samples, 3}, // shape
+                   {3 * byte_size, byte_size},           // C-style contiguous strides
+                   maps,                                 // the data pointer
+                   free_when_done);                      // numpy array references
+}
+
+py::array build_mapping(const py::array_t<int64_t> &docs_,
+                        const py::array_t<int> &sizes_,
+                        const int num_epochs,
+                        const uint64_t max_num_samples,
+                        const int max_seq_length,
+                        const double short_seq_prob,
+                        const int seed,
+                        const bool verbose,
+                        const int32_t min_num_sent)
+{
+
+  if (sizes_.size() > std::numeric_limits<uint32_t>::max())
+  {
+    if (verbose)
+    {
+      cout << "    using uint64 for data mapping..." << endl
+           << std::flush;
+    }
+    return build_mapping_impl<uint64_t>(docs_, sizes_, num_epochs,
+                                        max_num_samples, max_seq_length,
+                                        short_seq_prob, seed, verbose,
+                                        min_num_sent);
+  }
+  else
+  {
+    if (verbose)
+    {
+      cout << "    using uint32 for data mapping..." << endl
+           << std::flush;
+    }
+    return build_mapping_impl<uint32_t>(docs_, sizes_, num_epochs,
+                                        max_num_samples, max_seq_length,
+                                        short_seq_prob, seed, verbose,
+                                        min_num_sent);
+  }
+}
+
+template <typename DocIdx>
+py::array build_blocks_mapping_impl(const py::array_t<int64_t> &docs_,
+                                    const py::array_t<int32_t> &sizes_,
+                                    const py::array_t<int32_t> &titles_sizes_,
+                                    const int32_t num_epochs,
+                                    const uint64_t max_num_samples,
+                                    const int32_t max_seq_length,
+                                    const int32_t seed,
+                                    const bool verbose,
+                                    const bool use_one_sent_blocks)
+{
+  /* Build a mapping of (start-index, end-index, sequence-length) where
+       start and end index are the indices of the sentences in the sample
+       and sequence-length is the target sequence length.
+    */
+
+  // Consistency checks.
+  assert(num_epochs > 0);
+  assert(max_seq_length > 1);
+  assert(seed > 0);
+
+  // Remove bound checks.
+  auto docs = docs_.unchecked<1>();
+  auto sizes = sizes_.unchecked<1>();
+  auto titles_sizes = titles_sizes_.unchecked<1>();
+
+  if (verbose)
+  {
+    const auto sent_start_index = docs[0];
+    const auto sent_end_index = docs[docs_.shape(0) - 1];
+    const auto num_sentences = sent_end_index - sent_start_index;
+    cout << "    using:" << endl
+         << std::flush;
+    cout << "     number of documents:            " << docs_.shape(0) - 1 << endl
+         << std::flush;
+    cout << "     sentences range:                [" << sent_start_index << ", " << sent_end_index << ")" << endl
+         << std::flush;
+    cout << "     total number of sentences:      " << num_sentences << endl
+         << std::flush;
+    cout << "     number of epochs:               " << num_epochs << endl
+         << std::flush;
+    cout << "     maximum number of samples:      " << max_num_samples << endl
+         << std::flush;
+    cout << "     maximum sequence length:        " << max_seq_length << endl
+         << std::flush;
+    cout << "     seed:                           " << seed << endl
+         << std::flush;
+  }
+
+  // Mapping and its length (1D).
+  int64_t num_samples = -1;
+  DocIdx *maps = NULL;
+
+  // Acceptable number of sentences per block.
+  int min_num_sent = 2;
+  if (use_one_sent_blocks)
+  {
+    min_num_sent = 1;
+  }
+
+  // Perform two iterations, in the first iteration get the size
+  // and allocate memory and in the second iteration populate the map.
+  bool second = false;
+  for (int32_t iteration = 0; iteration < 2; ++iteration)
+  {
+
+    // Set the flag on second iteration.
+    second = (iteration == 1);
+
+    // Current map index.
+    uint64_t map_index = 0;
+
+    uint64_t empty_docs = 0;
+    uint64_t one_sent_docs = 0;
+    uint64_t long_sent_docs = 0;
+    // For each epoch:
+    for (int32_t epoch = 0; epoch < num_epochs; ++epoch)
+    {
+      // assign every block a unique id
+      int32_t block_id = 0;
+
+      if (map_index >= max_num_samples)
+      {
+        if (verbose && (!second))
+        {
+          cout << "    reached " << max_num_samples << " samples after "
+               << epoch << " epochs ..." << endl
+               << std::flush;
+        }
+        break;
+      }
+      // For each document:
+      for (int32_t doc = 0; doc < (docs.shape(0) - 1); ++doc)
+      {
+
+        // Document sentences are in [sent_index_first, sent_index_last)
+        const auto sent_index_first = docs[doc];
+        const auto sent_index_last = docs[doc + 1];
+        const auto target_seq_len = max_seq_length - titles_sizes[doc];
+
+        // At the begining of the document previous index is the
+        // start index.
+        auto prev_start_index = sent_index_first;
+
+        // Remaining documents.
+        auto num_remain_sent = sent_index_last - sent_index_first;
+
+        // Some bookkeeping
+        if ((epoch == 0) && (!second))
+        {
+          if (num_remain_sent == 0)
+          {
+            ++empty_docs;
+          }
+          if (num_remain_sent == 1)
+          {
+            ++one_sent_docs;
+          }
+        }
+        // Detect documents with long sentences.
+        bool contains_long_sentence = false;
+        if (num_remain_sent >= min_num_sent)
+        {
+          for (auto sent_index = sent_index_first;
+               sent_index < sent_index_last; ++sent_index)
+          {
+            if (sizes[sent_index] > LONG_SENTENCE_LEN)
+            {
+              if ((epoch == 0) && (!second))
+              {
+                ++long_sent_docs;
+              }
+              contains_long_sentence = true;
+              break;
+            }
+          }
+        }
+        // If we have enough sentences and no long sentences.
+        if ((num_remain_sent >= min_num_sent) && (!contains_long_sentence))
+        {
+
+          // Set values.
+          auto seq_len = int32_t{0};
+          auto num_sent = int32_t{0};
+
+          // Loop through sentences.
+          for (auto sent_index = sent_index_first;
+               sent_index < sent_index_last; ++sent_index)
+          {
+
+            // Add the size and number of sentences.
+            seq_len += sizes[sent_index];
+            ++num_sent;
+            --num_remain_sent;
+
+            // If we have reached the target length.
+            // and there are an acceptable number of sentences left
+            // and if we have at least the minimum number of sentences.
+            // or if we have reached end of the document.
+            if (((seq_len >= target_seq_len) &&
+                 (num_remain_sent >= min_num_sent) &&
+                 (num_sent >= min_num_sent)) ||
+                (num_remain_sent == 0))
+            {
+
+              // Populate the map.
+              if (second)
+              {
+                const auto map_index_0 = 4 * map_index;
+                // Each sample has 4 items: the starting sentence index, ending sentence index,
+                // the index of the document from which the block comes (used for fetching titles)
+                // and the unique id of the block (used for creating block indexes)
+
+                maps[map_index_0] = static_cast<DocIdx>(prev_start_index);
+                maps[map_index_0 + 1] = static_cast<DocIdx>(sent_index + 1);
+                maps[map_index_0 + 2] = static_cast<DocIdx>(doc);
+                maps[map_index_0 + 3] = static_cast<DocIdx>(block_id);
+              }
+
+              // Update indices / counters.
+              ++map_index;
+              ++block_id;
+              prev_start_index = sent_index + 1;
+              seq_len = 0;
+              num_sent = 0;
+            }
+          } // for (auto sent_index=sent_index_first; ...
+        }   // if (num_remain_sent > 1) {
+      }     // for (int doc=0; doc < num_docs; ++doc) {
+    }       // for (int epoch=0; epoch < num_epochs; ++epoch) {
+
+    if (!second)
+    {
+      if (verbose)
+      {
+        cout << "   number of empty documents: " << empty_docs << endl
+             << std::flush;
+        cout << "   number of documents with one sentence: " << one_sent_docs << endl
+             << std::flush;
+        cout << "   number of documents with long sentences: " << long_sent_docs << endl
+             << std::flush;
+        cout << "   will create mapping for " << map_index << " samples" << endl
+             << std::flush;
+      }
+      assert(maps == NULL);
+      assert(num_samples < 0);
+      maps = new DocIdx[4 * map_index];
+      num_samples = static_cast<int64_t>(map_index);
+    }
+
+  } // for (int iteration=0; iteration < 2; ++iteration) {
+
+  // Shuffle.
+  // We need a 64 bit random number generator as we might have more
+  // than 2 billion samples.
+  std::mt19937_64 rand64_gen(seed + 1);
+  for (auto i = (num_samples - 1); i > 0; --i)
+  {
+    const auto j = static_cast<int64_t>(rand64_gen() % (i + 1));
+    const auto i0 = 4 * i;
+    const auto j0 = 4 * j;
+    // Swap values.
+    swap(maps[i0], maps[j0]);
+    swap(maps[i0 + 1], maps[j0 + 1]);
+    swap(maps[i0 + 2], maps[j0 + 2]);
+    swap(maps[i0 + 3], maps[j0 + 3]);
+  }
+
+  // Method to deallocate memory.
+  py::capsule free_when_done(maps, [](void *mem_)
+                             {
+                               DocIdx *mem = reinterpret_cast<DocIdx *>(mem_);
+                               delete[] mem;
+                             });
+
+  // Return the numpy array.
+  const auto byte_size = sizeof(DocIdx);
+  return py::array(std::vector<int64_t>{num_samples, 4}, // shape
+                   {4 * byte_size, byte_size},           // C-style contiguous strides
+                   maps,                                 // the data pointer
+                   free_when_done);                      // numpy array references
+}
+
+py::array build_blocks_mapping(const py::array_t<int64_t> &docs_,
+                               const py::array_t<int> &sizes_,
+                               const py::array_t<int> &titles_sizes_,
+                               const int num_epochs,
+                               const uint64_t max_num_samples,
+                               const int max_seq_length,
+                               const int seed,
+                               const bool verbose,
+                               const bool use_one_sent_blocks)
+{
+
+  if (sizes_.size() > std::numeric_limits<uint32_t>::max())
+  {
+    if (verbose)
+    {
+      cout << "    using uint64 for data mapping..." << endl
+           << std::flush;
+    }
+    return build_blocks_mapping_impl<uint64_t>(docs_, sizes_, titles_sizes_,
+                                               num_epochs, max_num_samples, max_seq_length, seed, verbose, use_one_sent_blocks);
+  }
+  else
+  {
+    if (verbose)
+    {
+      cout << "    using uint32 for data mapping..." << endl
+           << std::flush;
+    }
+    return build_blocks_mapping_impl<uint32_t>(docs_, sizes_, titles_sizes_,
+                                               num_epochs, max_num_samples, max_seq_length, seed, verbose, use_one_sent_blocks);
+  }
+}
+
+PYBIND11_MODULE(helpers, m)
+{
+  m.def("build_mapping", &build_mapping);
+  m.def("build_blocks_mapping", &build_blocks_mapping);
+  m.def("build_sample_idx", &build_sample_idx);
+  m.def("build_blending_indices", &build_blending_indices);
+}

fengshen/data/megatron_dataloader/indexed_dataset.py

@@ -0,0 +1,585 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+# copied from fairseq/fairseq/data/indexed_dataset.py
+# Removed IndexedRawTextDataset since it relied on Fairseq dictionary
+# other slight modifications to remove fairseq dependencies
+# Added document index to index file and made it accessible.
+#    An empty sentence no longer separates documents.
+
+from functools import lru_cache
+import os
+import shutil
+import struct
+from itertools import accumulate
+
+import numpy as np
+import torch
+from fengshen.data.megatron_dataloader.utils import print_rank_0
+
+
+def __best_fitting_dtype(vocab_size=None):
+    if vocab_size is not None and vocab_size < 65500:
+        return np.uint16
+    else:
+        return np.int32
+
+
+def get_available_dataset_impl():
+    return ['lazy', 'cached', 'mmap']
+
+
+def infer_dataset_impl(path):
+    if IndexedDataset.exists(path):
+        with open(index_file_path(path), 'rb') as f:
+            magic = f.read(8)
+            if magic == IndexedDataset._HDR_MAGIC:
+                return 'cached'
+            elif magic == MMapIndexedDataset.Index._HDR_MAGIC[:8]:
+                return 'mmap'
+            else:
+                return None
+    else:
+        print(f"Dataset does not exist: {path}")
+        print("Path should be a basename that both .idx and "
+              ".bin can be appended to get full filenames.")
+        return None
+
+
+def make_builder(out_file, impl, vocab_size=None):
+    if impl == 'mmap':
+        return MMapIndexedDatasetBuilder(out_file,
+                                         dtype=__best_fitting_dtype(vocab_size))
+    else:
+        return IndexedDatasetBuilder(out_file)
+
+
+def make_dataset(path, impl, skip_warmup=False):
+    if not IndexedDataset.exists(path):
+        print(f"Dataset does not exist: {path}")
+        print("Path should be a basename that both .idx "
+              "and .bin can be appended to get full filenames.")
+        return None
+    if impl == 'infer':
+        impl = infer_dataset_impl(path)
+    if impl == 'lazy' and IndexedDataset.exists(path):
+        return IndexedDataset(path)
+    elif impl == 'cached' and IndexedDataset.exists(path):
+        return IndexedCachedDataset(path)
+    elif impl == 'mmap' and MMapIndexedDataset.exists(path):
+        return MMapIndexedDataset(path, skip_warmup)
+    print(f"Unknown dataset implementation: {impl}")
+    return None
+
+
+def dataset_exists(path, impl):
+    if impl == 'mmap':
+        return MMapIndexedDataset.exists(path)
+    else:
+        return IndexedDataset.exists(path)
+
+
+def read_longs(f, n):
+    a = np.empty(n, dtype=np.int64)
+    f.readinto(a)
+    return a
+
+
+def write_longs(f, a):
+    f.write(np.array(a, dtype=np.int64))
+
+
+dtypes = {
+    1: np.uint8,
+    2: np.int8,
+    3: np.int16,
+    4: np.int32,
+    5: np.int64,
+    6: np.float,
+    7: np.double,
+    8: np.uint16
+}
+
+
+def code(dtype):
+    for k in dtypes.keys():
+        if dtypes[k] == dtype:
+            return k
+    raise ValueError(dtype)
+
+
+def index_file_path(prefix_path):
+    return prefix_path + '.idx'
+
+
+def data_file_path(prefix_path):
+    return prefix_path + '.bin'
+
+
+def create_doc_idx(sizes):
+    doc_idx = [0]
+    for i, s in enumerate(sizes):
+        if s == 0:
+            doc_idx.append(i + 1)
+    return doc_idx
+
+
+class IndexedDataset(torch.utils.data.Dataset):
+    """Loader for IndexedDataset"""
+    _HDR_MAGIC = b'TNTIDX\x00\x00'
+
+    def __init__(self, path):
+        super().__init__()
+        self.path = path
+        self.data_file = None
+        self.read_index(path)
+
+    def read_index(self, path):
+        with open(index_file_path(path), 'rb') as f:
+            magic = f.read(8)
+            assert magic == self._HDR_MAGIC, (
+                'Index file doesn\'t match expected format. '
+                'Make sure that --dataset-impl is configured properly.'
+            )
+            version = f.read(8)
+            assert struct.unpack('<Q', version) == (1,)
+            code, self.element_size = struct.unpack('<QQ', f.read(16))
+            self.dtype = dtypes[code]
+            self._len, self.s = struct.unpack('<QQ', f.read(16))
+            self.doc_count = struct.unpack('<Q', f.read(8))
+            self.dim_offsets = read_longs(f, self._len + 1)
+            self.data_offsets = read_longs(f, self._len + 1)
+            self.sizes = read_longs(f, self.s)
+            self.doc_idx = read_longs(f, self.doc_count)
+
+    def read_data(self, path):
+        self.data_file = open(data_file_path(path), 'rb', buffering=0)
+
+    def check_index(self, i):
+        if i < 0 or i >= self._len:
+            raise IndexError('index out of range')
+
+    def __del__(self):
+        if self.data_file:
+            self.data_file.close()
+
+    # @lru_cache(maxsize=8)
+    def __getitem__(self, idx):
+        if not self.data_file:
+            self.read_data(self.path)
+        if isinstance(idx, int):
+            i = idx
+            self.check_index(i)
+            tensor_size = self.sizes[
+                self.dim_offsets[i]:self.dim_offsets[i + 1]]
+            a = np.empty(tensor_size, dtype=self.dtype)
+            self.data_file.seek(self.data_offsets[i] * self.element_size)
+            self.data_file.readinto(a)
+            return a
+        elif isinstance(idx, slice):
+            start, stop, step = idx.indices(len(self))
+            if step != 1:
+                raise ValueError(
+                    "Slices into indexed_dataset must be contiguous")
+            sizes = self.sizes[self.dim_offsets[start]:self.dim_offsets[stop]]
+            size = sum(sizes)
+            a = np.empty(size, dtype=self.dtype)
+            self.data_file.seek(self.data_offsets[start] * self.element_size)
+            self.data_file.readinto(a)
+            offsets = list(accumulate(sizes))
+            sents = np.split(a, offsets[:-1])
+            return sents
+
+    def __len__(self):
+        return self._len
+
+    def num_tokens(self, index):
+        return self.sizes[index]
+
+    def size(self, index):
+        return self.sizes[index]
+
+    @staticmethod
+    def exists(path):
+        return (
+            os.path.exists(index_file_path(path)) and os.path.exists(
+                data_file_path(path))
+        )
+
+    @property
+    def supports_prefetch(self):
+        return False  # avoid prefetching to save memory
+
+
+class IndexedCachedDataset(IndexedDataset):
+
+    def __init__(self, path):
+        super().__init__(path)
+        self.cache = None
+        self.cache_index = {}
+
+    @property
+    def supports_prefetch(self):
+        return True
+
+    def prefetch(self, indices):
+        if all(i in self.cache_index for i in indices):
+            return
+        if not self.data_file:
+            self.read_data(self.path)
+        indices = sorted(set(indices))
+        total_size = 0
+        for i in indices:
+            total_size += self.data_offsets[i + 1] - self.data_offsets[i]
+        self.cache = np.empty(total_size, dtype=self.dtype)
+        ptx = 0
+        self.cache_index.clear()
+        for i in indices:
+            self.cache_index[i] = ptx
+            size = self.data_offsets[i + 1] - self.data_offsets[i]
+            a = self.cache[ptx: ptx + size]
+            self.data_file.seek(self.data_offsets[i] * self.element_size)
+            self.data_file.readinto(a)
+            ptx += size
+        if self.data_file:
+            # close and delete data file after prefetch so we can pickle
+            self.data_file.close()
+            self.data_file = None
+
+    # @lru_cache(maxsize=8)
+    def __getitem__(self, idx):
+        if isinstance(idx, int):
+            i = idx
+            self.check_index(i)
+            tensor_size = self.sizes[
+                self.dim_offsets[i]:self.dim_offsets[i + 1]]
+            a = np.empty(tensor_size, dtype=self.dtype)
+            ptx = self.cache_index[i]
+            np.copyto(a, self.cache[ptx: ptx + a.size])
+            return a
+        elif isinstance(idx, slice):
+            # Hack just to make this work, can optimizer later if necessary
+            sents = []
+            for i in range(*idx.indices(len(self))):
+                sents.append(self[i])
+            return sents
+
+
+class IndexedDatasetBuilder(object):
+    element_sizes = {
+        np.uint8: 1,
+        np.int8: 1,
+        np.int16: 2,
+        np.int32: 4,
+        np.int64: 8,
+        np.float: 4,
+        np.double: 8
+    }
+
+    def __init__(self, out_file, dtype=np.int32):
+        self.out_file = open(out_file, 'wb')
+        self.dtype = dtype
+        self.data_offsets = [0]
+        self.dim_offsets = [0]
+        self.sizes = []
+        self.element_size = self.element_sizes[self.dtype]
+        self.doc_idx = [0]
+
+    def add_item(self, tensor):
+        bytes = self.out_file.write(np.array(tensor.numpy(), dtype=self.dtype))
+        self.data_offsets.append(
+            self.data_offsets[-1] + bytes / self.element_size)
+        for s in tensor.size():
+            self.sizes.append(s)
+        self.dim_offsets.append(self.dim_offsets[-1] + len(tensor.size()))
+
+    def end_document(self):
+        self.doc_idx.append(len(self.sizes))
+
+    def merge_file_(self, another_file):
+        index = IndexedDataset(another_file)
+        assert index.dtype == self.dtype
+
+        begin = self.data_offsets[-1]
+        for offset in index.data_offsets[1:]:
+            self.data_offsets.append(begin + offset)
+        self.sizes.extend(index.sizes)
+        begin = self.dim_offsets[-1]
+        for dim_offset in index.dim_offsets[1:]:
+            self.dim_offsets.append(begin + dim_offset)
+
+        with open(data_file_path(another_file), 'rb') as f:
+            while True:
+                data = f.read(1024)
+                if data:
+                    self.out_file.write(data)
+                else:
+                    break
+
+    def finalize(self, index_file):
+        self.out_file.close()
+        index = open(index_file, 'wb')
+        index.write(b'TNTIDX\x00\x00')
+        index.write(struct.pack('<Q', 1))
+        index.write(struct.pack('<QQ', code(self.dtype), self.element_size))
+        index.write(struct.pack('<QQ', len(
+            self.data_offsets) - 1, len(self.sizes)))
+        index.write(struct.pack('<Q', len(self.doc_idx)))
+        write_longs(index, self.dim_offsets)
+        write_longs(index, self.data_offsets)
+        write_longs(index, self.sizes)
+        write_longs(index, self.doc_idx)
+        index.close()
+
+
+def _warmup_mmap_file(path):
+    with open(path, 'rb') as stream:
+        while stream.read(100 * 1024 * 1024):
+            pass
+
+
+class MMapIndexedDataset(torch.utils.data.Dataset):
+    class Index(object):
+        _HDR_MAGIC = b'MMIDIDX\x00\x00'
+
+        @classmethod
+        def writer(cls, path, dtype):
+            class _Writer(object):
+                def __enter__(self):
+                    self._file = open(path, 'wb')
+
+                    self._file.write(cls._HDR_MAGIC)
+                    self._file.write(struct.pack('<Q', 1))
+                    self._file.write(struct.pack('<B', code(dtype)))
+
+                    return self
+
+                @staticmethod
+                def _get_pointers(sizes):
+                    dtype_size = dtype().itemsize
+                    address = 0
+                    pointers = []
+
+                    for size in sizes:
+                        pointers.append(address)
+                        address += size * dtype_size
+
+                    return pointers
+
+                def write(self, sizes, doc_idx):
+                    pointers = self._get_pointers(sizes)
+
+                    self._file.write(struct.pack('<Q', len(sizes)))
+                    self._file.write(struct.pack('<Q', len(doc_idx)))
+
+                    sizes = np.array(sizes, dtype=np.int32)
+                    self._file.write(sizes.tobytes(order='C'))
+                    del sizes
+
+                    pointers = np.array(pointers, dtype=np.int64)
+                    self._file.write(pointers.tobytes(order='C'))
+                    del pointers
+
+                    doc_idx = np.array(doc_idx, dtype=np.int64)
+                    self._file.write(doc_idx.tobytes(order='C'))
+
+                def __exit__(self, exc_type, exc_val, exc_tb):
+                    self._file.close()
+
+            return _Writer()
+
+        def __init__(self, path, skip_warmup=False):
+            with open(path, 'rb') as stream:
+                magic_test = stream.read(9)
+                assert self._HDR_MAGIC == magic_test, (
+                    'Index file doesn\'t match expected format. '
+                    'Make sure that --dataset-impl is configured properly.'
+                )
+                version = struct.unpack('<Q', stream.read(8))
+                assert (1,) == version
+
+                dtype_code, = struct.unpack('<B', stream.read(1))
+                self._dtype = dtypes[dtype_code]
+                self._dtype_size = self._dtype().itemsize
+
+                self._len = struct.unpack('<Q', stream.read(8))[0]
+                self._doc_count = struct.unpack('<Q', stream.read(8))[0]
+                offset = stream.tell()
+
+            if not skip_warmup:
+                print_rank_0("    warming up index mmap file...")
+                _warmup_mmap_file(path)
+
+            self._bin_buffer_mmap = np.memmap(path, mode='r', order='C')
+            self._bin_buffer = memoryview(self._bin_buffer_mmap)
+            print_rank_0("    reading sizes...")
+            self._sizes = np.frombuffer(
+                self._bin_buffer,
+                dtype=np.int32,
+                count=self._len,
+                offset=offset)
+            print_rank_0("    reading pointers...")
+            self._pointers = np.frombuffer(self._bin_buffer,
+                                           dtype=np.int64, count=self._len,
+                                           offset=offset + self._sizes.nbytes)
+            print_rank_0("    reading document index...")
+            self._doc_idx = np.frombuffer(
+                self._bin_buffer,
+                dtype=np.int64, count=self._doc_count,
+                offset=offset + self._sizes.nbytes + self._pointers.nbytes)
+
+        def __del__(self):
+            self._bin_buffer_mmap._mmap.close()
+            del self._bin_buffer_mmap
+
+        @property
+        def dtype(self):
+            return self._dtype
+
+        @property
+        def sizes(self):
+            return self._sizes
+
+        @property
+        def doc_idx(self):
+            return self._doc_idx
+
+        @lru_cache(maxsize=8)
+        def __getitem__(self, i):
+            return self._pointers[i], self._sizes[i]
+
+        def __len__(self):
+            return self._len
+
+    def __init__(self, path, skip_warmup=False):
+        super().__init__()
+
+        self._path = None
+        self._index = None
+        self._bin_buffer = None
+
+        self._do_init(path, skip_warmup)
+
+    def __getstate__(self):
+        return self._path
+
+    def __setstate__(self, state):
+        self._do_init(state)
+
+    def _do_init(self, path, skip_warmup):
+        self._path = path
+        self._index = self.Index(index_file_path(self._path), skip_warmup)
+
+        if not skip_warmup:
+            print_rank_0("    warming up data mmap file...")
+            _warmup_mmap_file(data_file_path(self._path))
+        print_rank_0("    creating numpy buffer of mmap...")
+        self._bin_buffer_mmap = np.memmap(
+            data_file_path(self._path), mode='r', order='C')
+        print_rank_0("    creating memory view of numpy buffer...")
+        self._bin_buffer = memoryview(self._bin_buffer_mmap)
+
+    def __del__(self):
+        self._bin_buffer_mmap._mmap.close()
+        del self._bin_buffer_mmap
+        del self._index
+
+    def __len__(self):
+        return len(self._index)
+
+    # @lru_cache(maxsize=8)
+    def __getitem__(self, idx):
+        if isinstance(idx, int):
+            ptr, size = self._index[idx]
+            np_array = np.frombuffer(self._bin_buffer, dtype=self._index.dtype,
+                                     count=size, offset=ptr)
+            return np_array
+        elif isinstance(idx, slice):
+            start, stop, step = idx.indices(len(self))
+            if step != 1:
+                raise ValueError(
+                    "Slices into indexed_dataset must be contiguous")
+            ptr = self._index._pointers[start]
+            sizes = self._index._sizes[idx]
+            offsets = list(accumulate(sizes))
+            total_size = sum(sizes)
+            np_array = np.frombuffer(self._bin_buffer, dtype=self._index.dtype,
+                                     count=total_size, offset=ptr)
+            sents = np.split(np_array, offsets[:-1])
+            return sents
+
+    def get(self, idx, offset=0, length=None):
+        """ Retrieves a single item from the dataset with the option to only
+        return a portion of the item.
+
+        get(idx) is the same as [idx] but get() does not support slicing.
+        """
+        ptr, size = self._index[idx]
+        if length is None:
+            length = size - offset
+        ptr += offset * np.dtype(self._index.dtype).itemsize
+        np_array = np.frombuffer(self._bin_buffer, dtype=self._index.dtype,
+                                 count=length, offset=ptr)
+        return np_array
+
+    @property
+    def sizes(self):
+        return self._index.sizes
+
+    @property
+    def doc_idx(self):
+        return self._index.doc_idx
+
+    def get_doc_idx(self):
+        return self._index._doc_idx
+
+    def set_doc_idx(self, doc_idx_):
+        self._index._doc_idx = doc_idx_
+
+    @property
+    def supports_prefetch(self):
+        return False
+
+    @staticmethod
+    def exists(path):
+        return (
+            os.path.exists(index_file_path(path)) and os.path.exists(
+                data_file_path(path))
+        )
+
+
+class MMapIndexedDatasetBuilder(object):
+    def __init__(self, out_file, dtype=np.int64):
+        self._data_file = open(out_file, 'wb', buffering=5000000)
+        self._dtype = dtype
+        self._sizes = []
+        self._doc_idx = [0]
+
+    def add_item(self, tensor):
+        np_array = np.array(tensor.numpy(), dtype=self._dtype)
+        self._data_file.write(np_array.tobytes(order='C'))
+        self._sizes.append(np_array.size)
+
+    def end_document(self):
+        self._doc_idx.append(len(self._sizes))
+
+    def merge_file_(self, another_file):
+        # Concatenate index
+        index = MMapIndexedDataset.Index(index_file_path(another_file))
+        assert index.dtype == self._dtype
+
+        for size in index.sizes:
+            self._sizes.append(size)
+
+        # Concatenate data
+        with open(data_file_path(another_file), 'rb') as f:
+            shutil.copyfileobj(f, self._data_file)
+
+    def finalize(self, index_file):
+        self._data_file.close()
+
+        with MMapIndexedDataset.Index.writer(index_file, self._dtype) as index:
+            index.write(self._sizes, self._doc_idx)

fengshen/data/megatron_dataloader/utils.py

@@ -0,0 +1,24 @@
+# coding=utf-8
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import torch
+
+
+def print_rank_0(message):
+    """If distributed is initialized, print only on rank 0."""
+    if torch.distributed.is_initialized():
+        if torch.distributed.get_rank() == 0:
+            print(message, flush=True)
+    else:
+        print(message, flush=True)

fengshen/data/mmap_dataloader/mmap_datamodule.py

@@ -0,0 +1,68 @@
+from typing import Optional
+from pytorch_lightning import LightningDataModule
+from torch.utils.data import DataLoader
+from fengshen.data.mmap_index_dataset import MMapIndexDataset
+
+
+class MMapDataModule(LightningDataModule):
+    @ staticmethod
+    def add_data_specific_args(parent_args):
+        parser = parent_args.add_argument_group('MMAP DataModule')
+        parser.add_argument('--num_workers', default=8, type=int)
+        parser.add_argument('--train_batchsize', default=32, type=int)
+        parser.add_argument('--eval_batchsize', default=32, type=int)
+        parser.add_argument('--test_batchsize', default=32, type=int)
+        parser.add_argument('--train_datas', default=[
+            './train_datas'
+        ], type=str, nargs='+')
+        parser.add_argument('--valid_datas', default=[
+            './valid_datas'
+        ], type=str, nargs='+')
+        parser.add_argument('--test_datas', default=[
+            './test_datas'],
+            type=str, nargs='+')
+        parser.add_argument('--input_tensor_name', default=['input_ids'], type=str, nargs='+')
+        return parent_args
+
+    def __init__(
+        self,
+        collate_fn,
+        args,
+        **kwargs,
+    ):
+        super().__init__()
+        self.collate_fn = collate_fn
+        self.train_dataset = MMapIndexDataset(args.train_datas, args.input_tensor_name)
+        self.valid_dataset = MMapIndexDataset(args.valid_datas, args.input_tensor_name)
+        self.test_dataset = MMapIndexDataset(args.test_datas, args.input_tensor_name)
+        self.save_hyperparameters(args)
+
+    def setup(self, stage: Optional[str] = None) -> None:
+        return super().setup(stage)
+
+    def train_dataloader(self):
+        return DataLoader(
+            self.train_dataset,
+            batch_size=self.hparams.train_batchsize,
+            shuffle=True,
+            num_workers=self.hparams.num_workers,
+            collate_fn=self.collate_fn,
+        )
+
+    def val_dataloader(self):
+        return DataLoader(
+            self.valid_dataset,
+            batch_size=self.hparams.eval_batchsize,
+            shuffle=True,
+            num_workers=self.hparams.num_workers,
+            collate_fn=self.collate_fn,
+        )
+
+    def test_dataloader(self):
+        return DataLoader(
+            self.test_dataset,
+            batch_size=self.hparams.test_batchsize,
+            shuffle=True,
+            num_workers=self.hparams.num_workers,
+            collate_fn=self.collate_fn,
+        )

fengshen/data/mmap_dataloader/mmap_index_dataset.py

@@ -0,0 +1,53 @@
+import numpy as np
+import torch
+from typing import List
+from torch.utils.data import Dataset
+
+
+class MMapIndexDataset(Dataset):
+    # datapaths 是所有的内存映射文件的路径
+    # input_tensor_name 是输入的tensor的名字 例如 ['input_ids'] 会存储在对应的文件里面
+    def __init__(self, datapaths: List[str], input_tensor_name: List[str]):
+        dict_idx_fp = {}
+        dict_bin_fp = {}
+        idx_len = []
+        for tensor_name in input_tensor_name:
+            idx_fp = []
+            bin_fp = []
+            len = 0
+            for data_path in datapaths:
+                idx_fp += [np.load(
+                    data_path + '_' + tensor_name + '.npy', mmap_mode='r')]
+                bin_fp += [np.memmap(
+                    data_path + '_' + tensor_name + '.bin',
+                    dtype='long',
+                    mode='r')]
+                len += idx_fp[-1].shape[0]
+                idx_len += [idx_fp[-1].shape[0]]
+            dict_idx_fp[tensor_name] = idx_fp
+            dict_bin_fp[tensor_name] = bin_fp
+            #  通常情况下不同的tensor的长度是一样的
+            self._len = len
+
+        self._input_tensor_name = input_tensor_name
+        self._dict_idx_fp = dict_idx_fp
+        self._dict_bin_fp = dict_bin_fp
+        self._idx_len = idx_len
+
+    def __len__(self):
+        return self._len
+
+    def __getitem__(self, idx):
+        sample = {}
+        for i in range(len(self._idx_len)):
+            if idx >= self._idx_len[i]:
+                idx -= self._idx_len[i]
+            else:
+                break
+        for tensor_name in self._input_tensor_name:
+            sample[tensor_name] = torch.tensor(self._dict_bin_fp[tensor_name][i][
+                self._dict_idx_fp[tensor_name][i][idx, 0]:
+                    self._dict_idx_fp[tensor_name][i][idx, 1]
+            ], dtype=torch.long)
+        # print(sample)
+        return sample

fengshen/data/preprocess.py

@@ -0,0 +1 @@
+# coding=utf-8

fengshen/data/sequence_tagging_dataloader/sequence_tagging_collator.py

@@ -0,0 +1,274 @@
+from dataclasses import dataclass
+from torch.utils.data._utils.collate import default_collate
+
+import copy
+import torch
+import numpy as np
+
+@dataclass
+class CollatorForLinear:
+    args = None
+    tokenizer = None
+    label2id = None
+
+    def __call__(self, samples):
+        cls_token = "[CLS]"
+        sep_token = "[SEP]"
+        pad_token = 0
+        special_tokens_count = 2
+        segment_id = 0
+
+        features=[]
+
+        for (ex_index, example) in enumerate(samples):
+            tokens = copy.deepcopy(example['text_a'])
+
+            label_ids = [self.label2id[x] for x in example['labels']]
+
+            if len(tokens) > self.args.max_seq_length - special_tokens_count:
+                tokens = tokens[: (self.args.max_seq_length - special_tokens_count)]
+                label_ids = label_ids[: (self.args.max_seq_length - special_tokens_count)]
+
+            tokens += [sep_token]
+            label_ids += [self.label2id["O"]]
+            segment_ids = [segment_id] * len(tokens)
+
+            tokens = [cls_token] + tokens
+            label_ids = [self.label2id["O"]] + label_ids
+            segment_ids = [segment_id] + segment_ids
+
+            input_ids = self.tokenizer.convert_tokens_to_ids(tokens)
+            input_mask = [1] * len(input_ids)
+            input_len = len(label_ids)
+            padding_length = self.args.max_seq_length - len(input_ids)
+
+            input_ids += [pad_token] * padding_length
+            input_mask += [0] * padding_length
+            segment_ids += [segment_id] * padding_length
+            label_ids += [pad_token] * padding_length
+
+            assert len(input_ids) == self.args.max_seq_length
+            assert len(input_mask) == self.args.max_seq_length
+            assert len(segment_ids) == self.args.max_seq_length
+            assert len(label_ids) == self.args.max_seq_length
+
+            features.append({
+                    'input_ids':torch.tensor(input_ids),
+                    'attention_mask':torch.tensor(input_mask),
+                    'input_len':torch.tensor(input_len),
+                    'token_type_ids':torch.tensor(segment_ids),
+                    'labels':torch.tensor(label_ids),
+            })
+
+        return default_collate(features)
+
+@dataclass
+class CollatorForCrf:
+    args = None
+    tokenizer = None
+    label2id = None
+
+    def __call__(self, samples):
+        features = []
+        cls_token = "[CLS]"
+        sep_token = "[SEP]"
+        pad_token = 0
+        special_tokens_count = 2
+        segment_id = 0
+
+        for (ex_index, example) in enumerate(samples):
+            tokens = copy.deepcopy(example['text_a'])
+
+            label_ids = [self.label2id[x] for x in example['labels']]
+
+            if len(tokens) > self.args.max_seq_length - special_tokens_count:
+                tokens = tokens[: (self.args.max_seq_length - special_tokens_count)]
+                label_ids = label_ids[: (self.args.max_seq_length - special_tokens_count)]
+
+            tokens += [sep_token]
+            label_ids += [self.label2id["O"]]
+            segment_ids = [segment_id] * len(tokens)
+
+            tokens = [cls_token] + tokens
+            label_ids = [self.label2id["O"]] + label_ids
+            segment_ids = [segment_id] + segment_ids
+
+            input_ids = self.tokenizer.convert_tokens_to_ids(tokens)
+            input_mask = [1] * len(input_ids)
+            input_len = len(label_ids)
+            padding_length = self.args.max_seq_length - len(input_ids)
+
+            input_ids += [pad_token] * padding_length
+            input_mask += [0] * padding_length
+            segment_ids += [segment_id] * padding_length
+            label_ids += [pad_token] * padding_length
+
+            assert len(input_ids) == self.args.max_seq_length
+            assert len(input_mask) == self.args.max_seq_length
+            assert len(segment_ids) == self.args.max_seq_length
+            assert len(label_ids) == self.args.max_seq_length
+
+            features.append({
+                    'input_ids':torch.tensor(input_ids),
+                    'attention_mask':torch.tensor(input_mask),
+                    'input_len':torch.tensor(input_len),
+                    'token_type_ids':torch.tensor(segment_ids),
+                    'labels':torch.tensor(label_ids),
+            })
+        
+        return default_collate(features)
+
+
+@dataclass
+class CollatorForSpan:
+    args = None
+    tokenizer = None
+    label2id = None
+
+    def __call__(self, samples):
+
+        features = []
+        cls_token = "[CLS]"
+        sep_token = "[SEP]"
+        pad_token = 0
+        special_tokens_count = 2
+        max_entities_count = 100
+        segment_id = 0
+
+        for (ex_index, example) in enumerate(samples):
+            subjects = copy.deepcopy(example['subject'])
+            tokens = copy.deepcopy(example['text_a'])
+            start_ids = [0] * len(tokens)
+            end_ids = [0] * len(tokens)
+            subject_ids = []
+            for subject in subjects:
+                label = subject[0]
+                start = subject[1]
+                end = subject[2]
+                start_ids[start] = self.label2id[label]
+                end_ids[end] = self.label2id[label]
+                subject_ids.append([self.label2id[label], start, end])
+            
+            subject_ids+=[[-1,-1,-1]]*(max_entities_count-len(subject_ids))
+
+            if len(tokens) > self.args.max_seq_length - special_tokens_count:
+                tokens = tokens[: (self.args.max_seq_length - special_tokens_count)]
+                start_ids = start_ids[: (self.args.max_seq_length - special_tokens_count)]
+                end_ids = end_ids[: (self.args.max_seq_length - special_tokens_count)]
+
+            tokens += [sep_token]
+            start_ids += [0]
+            end_ids += [0]
+            segment_ids = [segment_id] * len(tokens)
+
+            tokens = [cls_token] + tokens
+            start_ids = [0] + start_ids
+            end_ids = [0] + end_ids
+            segment_ids = [segment_id] + segment_ids
+
+            input_ids = self.tokenizer.convert_tokens_to_ids(tokens)
+            input_mask = [1] * len(input_ids)
+            input_len = len(input_ids)
+            padding_length = self.args.max_seq_length - len(input_ids)
+
+            input_ids += [pad_token] * padding_length
+            input_mask += [0] * padding_length
+            segment_ids += [segment_id] * padding_length
+            start_ids += [0] * padding_length
+            end_ids += [0] * padding_length
+
+            assert len(input_ids) == self.args.max_seq_length
+            assert len(input_mask) == self.args.max_seq_length
+            assert len(segment_ids) == self.args.max_seq_length
+            assert len(start_ids) == self.args.max_seq_length
+            assert len(end_ids) == self.args.max_seq_length
+
+            features.append({
+                    'input_ids': torch.tensor(np.array(input_ids)),
+                    'attention_mask': torch.tensor(np.array(input_mask)),
+                    'token_type_ids': torch.tensor(np.array(segment_ids)),
+                    'start_positions': torch.tensor(np.array(start_ids)),
+                    'end_positions': torch.tensor(np.array(end_ids)),
+                    "subjects": torch.tensor(np.array(subject_ids)),
+                    'input_len': torch.tensor(np.array(input_len)),
+                })
+        
+        return default_collate(features)
+
+
+@dataclass
+class CollatorForBiaffine:
+    args = None
+    tokenizer = None
+    label2id = None
+
+    
+    def __call__(self, samples):
+
+        features = []
+        cls_token = "[CLS]"
+        sep_token = "[SEP]"
+        pad_token = 0
+        special_tokens_count = 2
+        segment_id = 0
+
+        for (ex_index, example) in enumerate(samples):
+            subjects = copy.deepcopy(example['subject'])
+            tokens = copy.deepcopy(example['text_a'])
+
+            span_labels = np.zeros((self.args.max_seq_length,self.args.max_seq_length))
+            span_labels[:] = self.label2id["O"]
+
+            for subject in subjects:
+                label = subject[0]
+                start = subject[1]
+                end = subject[2]
+                if start < self.args.max_seq_length - special_tokens_count and end < self.args.max_seq_length - special_tokens_count:
+                    span_labels[start + 1, end + 1] = self.label2id[label]
+
+            if len(tokens) > self.args.max_seq_length - special_tokens_count:
+                tokens = tokens[: (self.args.max_seq_length - special_tokens_count)]
+
+            tokens += [sep_token]
+            span_labels[len(tokens), :] = self.label2id["O"]
+            span_labels[:, len(tokens)] = self.label2id["O"]
+            segment_ids = [segment_id] * len(tokens)
+
+            tokens = [cls_token] + tokens
+            span_labels[0, :] = self.label2id["O"]
+            span_labels[:, 0] = self.label2id["O"]
+            segment_ids = [segment_id] + segment_ids
+
+            input_ids = self.tokenizer.convert_tokens_to_ids(tokens)
+            input_mask = [0] * len(input_ids)
+            span_mask = np.ones(span_labels.shape)
+            input_len = len(input_ids)
+
+            padding_length = self.args.max_seq_length - len(input_ids)
+
+            input_ids += [pad_token] * padding_length
+            input_mask += [0] * padding_length
+            segment_ids += [segment_id] * padding_length
+            span_labels[input_len:, :] = 0
+            span_labels[:, input_len:] = 0
+            span_mask[input_len:, :] = 0
+            span_mask[:, input_len:] = 0
+            span_mask=np.triu(span_mask,0)
+            span_mask=np.tril(span_mask,10)
+
+            assert len(input_ids) == self.args.max_seq_length
+            assert len(input_mask) == self.args.max_seq_length
+            assert len(segment_ids) == self.args.max_seq_length
+            assert len(span_labels) == self.args.max_seq_length
+            assert len(span_labels[0]) == self.args.max_seq_length
+
+            features.append({
+                    'input_ids': torch.tensor(np.array(input_ids)),
+                    'attention_mask': torch.tensor(np.array(input_mask)),
+                    'token_type_ids': torch.tensor(np.array(segment_ids)),
+                    'span_labels': torch.tensor(np.array(span_labels)),
+                    'span_mask': torch.tensor(np.array(span_mask)),
+                    'input_len': torch.tensor(np.array(input_len)),
+            })
+        
+        return default_collate(features)

fengshen/data/sequence_tagging_dataloader/sequence_tagging_datasets.py

@@ -0,0 +1,116 @@
+from torch.utils.data import Dataset
+from fengshen.metric.utils_ner import get_entities
+
+import os
+
+def get_datasets(args):
+    processor = DataProcessor(args.data_dir, args.decode_type)
+
+    train_data = TaskDataset(processor=processor, mode="train")
+    valid_data = TaskDataset(processor=processor, mode="dev")
+    test_data = TaskDataset(processor=processor, mode="dev")
+
+    return {"train":train_data,"validation":valid_data,"test":test_data}
+
+# def get_labels(decode_type):
+#     with open("/cognitive_comp/lujunyu/data_zh/NER_Aligned/weibo/labels.txt") as f:
+#         label_list = ["[PAD]", "[START]", "[END]"]
+
+#         if decode_type=="crf" or decode_type=="linear":
+#             for line in f.readlines():
+#                 label_list.append(line.strip())
+#         elif decode_type=="biaffine" or decode_type=="span":
+#             for line in f.readlines():
+#                 tag = line.strip().split("-")
+#                 if len(tag) == 1 and tag[0] not in label_list:
+#                     label_list.append(tag[0])
+#                 elif tag[1] not in label_list:
+#                     label_list.append(tag[1])
+    
+#     label2id={label:id for id,label in enumerate(label_list)}
+#     id2label={id:label for id,label in enumerate(label_list)}
+#     return label2id, id2label
+
+class DataProcessor(object):
+    def __init__(self, data_dir, decode_type) -> None:
+        super().__init__()
+        self.data_dir = data_dir
+        self.decode_type = decode_type
+
+    def get_examples(self, mode):
+        return self._create_examples(self._read_text(os.path.join(self.data_dir, mode + ".all.bmes")), mode)
+
+    @staticmethod
+    def get_labels(args):
+        with open(os.path.join(args.data_dir, "labels.txt")) as f:
+            label_list = ["[PAD]", "[START]", "[END]"]
+
+            if args.decode_type=="crf" or args.decode_type=="linear":
+                for line in f.readlines():
+                    label_list.append(line.strip())
+            elif args.decode_type=="biaffine" or args.decode_type=="span":
+                for line in f.readlines():
+                    tag = line.strip().split("-")
+                    if len(tag) == 1 and tag[0] not in label_list:
+                        label_list.append(tag[0])
+                    elif tag[1] not in label_list:
+                        label_list.append(tag[1])
+
+        label2id = {label: i for i, label in enumerate(label_list)}
+        id2label={id:label for id,label in enumerate(label_list)}
+        return label2id,id2label
+
+    def _create_examples(self, lines, set_type):
+        examples = []
+        for (i, line) in enumerate(lines):
+            guid = "%s-%s" % (set_type, i)
+            text_a = line['words']
+            labels = []
+            for x in line['labels']:
+                if 'M-' in x:
+                    labels.append(x.replace('M-', 'I-'))
+                else:
+                    labels.append(x)
+            subject = get_entities(labels, id2label=None, markup='bioes')
+            examples.append({'guid':guid, 'text_a':text_a, 'labels':labels, 'subject':subject})
+        return examples
+
+    @classmethod
+    def _read_text(self, input_file):
+        lines = []
+        with open(input_file, 'r') as f:
+            words = []
+            labels = []
+            for line in f:
+                if line.startswith("-DOCSTART-") or line == "" or line == "\n":
+                    if words:
+                        lines.append({"words": words, "labels": labels})
+                        words = []
+                        labels = []
+                else:
+                    splits = line.split()
+                    words.append(splits[0])
+                    if len(splits) > 1:
+                        labels.append(splits[-1].replace("\n", ""))
+                    else:
+                        # Examples could have no label for mode = "test"
+                        labels.append("O")
+            if words:
+                lines.append({"words": words, "labels": labels})
+        return lines
+
+
+class TaskDataset(Dataset):
+    def __init__(self, processor, mode='train'):
+        super().__init__()
+        self.data = self.load_data(processor, mode)
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, index):
+        return self.data[index]
+
+    def load_data(self, processor, mode):
+        examples = processor.get_examples(mode)
+        return examples

fengshen/data/t5_dataloader/t5_datasets.py

@@ -0,0 +1,562 @@
+# coding=utf8
+import json
+from torch.utils.data import Dataset, DataLoader
+from tqdm import tqdm
+from transformers import BertTokenizer, MT5Config, MT5Tokenizer, BatchEncoding
+import torch
+import pytorch_lightning as pl
+import numpy as np
+from itertools import chain
+import sys
+sys.path.append('../../')
+
+
+def compute_input_and_target_lengths(inputs_length, noise_density, mean_noise_span_length):
+    """This function is copy of `random_spans_helper <https://github.com/google-research/
+    text-to-text-transfer-transformer/blob/84f8bcc14b5f2c03de51bd3587609ba8f6bbd1cd/t5/data/preprocessors.py#L2466>`__ .
+    Training parameters to avoid padding with random_spans_noise_mask.
+    When training a model with random_spans_noise_mask, we would like to set the other
+    training hyperparmeters in a way that avoids padding.
+    This function helps us compute these hyperparameters.
+    We assume that each noise span in the input is replaced by extra_tokens_per_span_inputs sentinel tokens,
+    and each non-noise span in the targets is replaced by extra_tokens_per_span_targets sentinel tokens.
+    This function tells us the required number of tokens in the raw example (for split_tokens())
+    as well as the length of the encoded targets. Note that this function assumes
+    the inputs and targets will have EOS appended and includes that in the reported length.
+    Args:
+        inputs_length: an integer - desired length of the tokenized inputs sequence
+        noise_density: a float
+        mean_noise_span_length: a float
+    Returns:
+        tokens_length: length of original text in tokens
+        targets_length: an integer - length in tokens of encoded targets sequence
+    """
+
+    def _tokens_length_to_inputs_length_targets_length(tokens_length):
+        num_noise_tokens = int(round(tokens_length * noise_density))
+        num_nonnoise_tokens = tokens_length - num_noise_tokens
+        num_noise_spans = int(round(num_noise_tokens / mean_noise_span_length))
+        # inputs contain all nonnoise tokens, sentinels for all noise spans
+        # and one EOS token.
+        _input_length = num_nonnoise_tokens + num_noise_spans + 1
+        _output_length = num_noise_tokens + num_noise_spans + 1
+        return _input_length, _output_length
+
+    tokens_length = inputs_length
+
+    while _tokens_length_to_inputs_length_targets_length(tokens_length + 1)[0] <= inputs_length:
+        tokens_length += 1
+
+    inputs_length, targets_length = _tokens_length_to_inputs_length_targets_length(
+        tokens_length)
+
+    # minor hack to get the targets length to be equal to inputs length
+    # which is more likely to have been set to a nice round number.
+    if noise_density == 0.5 and targets_length > inputs_length:
+        tokens_length -= 1
+        targets_length -= 1
+    return tokens_length, targets_length
+
+
+class UnsuperviseT5Dataset(Dataset):
+    '''
+    Dataset Used for T5 unsuprvise pretrain.
+    load_data_type = 0: load raw data from data path and save tokenized data, call function load_data
+    load_data_type = 1: load tokenized data from path, call function load_tokenized_data
+    load_data_type = 2: load tokenized data from memery data, call function load_tokenized_memory_data
+    '''
+
+    def __init__(self, data_path, args, load_data_type=0, data=None):
+        super().__init__()
+
+        if args.tokenizer_type == 't5_tokenizer':
+            if args.new_vocab_path is not None:
+                self.tokenizer = MT5Tokenizer.from_pretrained(args.new_vocab_path)
+            else:
+                self.tokenizer = MT5Tokenizer.from_pretrained(args.pretrained_model_path)
+        else:
+            self.tokenizer = BertTokenizer.from_pretrained(args.pretrained_model_path)
+        self.noise_density = 0.15
+        self.mean_noise_span_length = 3
+        self.text_column_name = args.text_column_name
+        self.dataset_num_workers = args.dataset_num_workers
+        self.max_seq_length = args.max_seq_length
+        self.remove_columns = args.remove_columns
+        # whether load tokenieze data
+        self.load_data_type = load_data_type
+
+        if self.load_data_type == 0:
+            # T5-like span masked language modeling will fuse consecutively masked tokens to a single sentinel token.
+            # To ensure that the input length is `max_seq_length`, we need to increase the maximum length
+            # according to `mlm_probability` and `mean_noise_span_length`.
+            # We can also define the label length accordingly.
+            self.expanded_inputs_length, self.targets_length = compute_input_and_target_lengths(
+                inputs_length=self.max_seq_length,
+                noise_density=self.noise_density,
+                mean_noise_span_length=self.mean_noise_span_length,
+            )
+            print('self.expanded_inputs_length, self.targets_length:{},{}'.format(
+                self.expanded_inputs_length, self.targets_length))
+            self.data = self.load_data(data_path)
+        elif self.load_data_type == 1:
+            self.data = self.load_tokenized_data(data_path)
+        else:
+            assert data is not None
+            self.data = self.load_tokenized_memory_data(data)
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, index):
+        return self.data[index]
+
+    def load_data(self, data_path):
+        # TODO: large data process
+        from data.fs_datasets import load_dataset
+        samples = load_dataset(
+            # samples = datasets.load_from_disk(data_path)['train']
+            data_path, num_proc=self.dataset_num_workers)['train']
+        # print(samples)
+        tokenized_datasets = samples.map(
+            self.tokenize_function,
+            batched=True,
+            num_proc=self.dataset_num_workers,
+            # load_from_cache_file=not data_args.overwrite_cache,
+        ).map(
+            batched=True,
+            num_proc=self.dataset_num_workers,
+            remove_columns=self.remove_columns)
+        # Note that with `batched=True`, this map processes 1,000 texts together, so group_texts throws away a
+        # remainder for each of those groups of 1,000 texts. You can adjust that batch_size here but a higher value
+        # might be slower to preprocess.
+        #
+        # To speed up this part, we use multiprocessing. See the documentation of the map method for more information:
+        # https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasets.Dataset.map
+        tokenized_datasets = tokenized_datasets.map(
+            self.group_texts,
+            batched=True,
+            num_proc=self.dataset_num_workers,
+            # load_from_cache_file=not data_args.overwrite_cache,
+        )
+        return tokenized_datasets
+    '''
+        The function load tokenized data saved from load_data function.
+    '''
+
+    def load_tokenized_data(self, data_path):
+        from data.fs_datasets import load_dataset
+        samples = load_dataset(data_path)['train']
+        return samples
+
+    def load_tokenized_memory_data(self, data):
+        return data
+
+    # Otherwise, we tokenize every text, then concatenate them together before splitting them in smaller parts.
+    # Since we make sure that all sequences are of the same length, no attention_mask is needed.
+    def tokenize_function(self, examples):
+        # 这里add_special_tokens=False，避免句子中间出现eos
+        return self.tokenizer(examples[self.text_column_name],
+                              add_special_tokens=False,
+                              return_attention_mask=False)
+
+    # Main data processing function that will concatenate all texts from our dataset
+    # and generate chunks of expanded_inputs_length.
+    def group_texts(self, examples):
+        # Concatenate all texts.
+        concatenated_examples = {
+            k: list(chain(*examples[k])) for k in examples.keys()}
+        total_length = len(concatenated_examples[list(examples.keys())[0]])
+        # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
+        # customize this part to your needs.
+        if total_length >= self.expanded_inputs_length:
+            total_length = (
+                total_length // self.expanded_inputs_length) * self.expanded_inputs_length
+        # Split by chunks of max_len.
+        result = {
+            k: [t[i: i + self.expanded_inputs_length]
+                for i in range(0, total_length, self.expanded_inputs_length)]
+            for k, t in concatenated_examples.items()
+        }
+        return result
+
+
+class UnsuperviseT5DataModel(pl.LightningDataModule):
+    @staticmethod
+    def add_data_specific_args(parent_args):
+        parser = parent_args.add_argument_group('UnsuperviseT5DataModel')
+        parser.add_argument('--dataset_num_workers', default=8, type=int)
+        parser.add_argument('--dataloader_num_workers', default=4, type=int)
+        parser.add_argument(
+            '--train_data_path', default='wudao_180g_mt5_tokenized', type=str)
+        parser.add_argument('--train_batchsize', default=2, type=int)
+        parser.add_argument('--valid_batchsize', default=2, type=int)
+        parser.add_argument('--train_split_size', default=None, type=float)
+        parser.add_argument('--tokenizer_type', default='t5_tokenizer', choices=['t5_tokenizer', 'bert_tokenizer'])
+        parser.add_argument('--text_column_name', default='text')
+        parser.add_argument('--remove_columns', nargs='+', default=[])
+        return parent_args
+
+    def __init__(self, args):
+        super().__init__()
+        self.save_hyperparameters(args)
+        if args.train_split_size is not None:
+            from data.fs_datasets import load_dataset
+            data_splits = load_dataset(args.train_data_path, num_proc=args.dataset_num_workers)
+            train_split = data_splits['train']
+            test_split = data_splits['test']
+            print('train:', train_split, '\ntest_data:', test_split)
+            self.train_dataset = UnsuperviseT5Dataset('', args, load_data_type=2, data=train_split)
+            self.test_dataset = UnsuperviseT5Dataset('', args, load_data_type=2, data=test_split)
+        else:
+            self.train_data = UnsuperviseT5Dataset(args.train_data_path, args, load_data_type=1)
+
+        self.config = MT5Config.from_pretrained(args.pretrained_model_path)
+        self.noise_density = 0.15
+        self.mean_noise_span_length = 3
+        self.pad_token_id = self.config.pad_token_id
+        self.decoder_start_token_id = self.config.decoder_start_token_id
+        self.eos_token_id = self.config.eos_token_id
+        self.vocab_size = self.config.vocab_size
+        self.max_seq_length = args.max_seq_length
+        # 因为加载旧的spm里面已经包括了exrta_ids，但是T5Tokenizer会在spm的基础上再增加100个extra_ids,所以需要指定extra_ids=0
+        if args.tokenizer_type == 't5_tokenizer' and args.new_vocab_path is not None:
+            self.tokenizer = MT5Tokenizer.from_pretrained(args.new_vocab_path, extra_ids=0)
+            # 如果是刚开始加载mt5,需要更新vocab_size为提取中英词之后的new_vocab_size
+            self.vocab_size = len(self.tokenizer)
+
+        # T5-like span masked language modeling will fuse consecutively masked tokens to a single sentinel token.
+        # To ensure that the input length is `max_seq_length`, we need to increase the maximum length
+        # according to `mlm_probability` and `mean_noise_span_length`. We can also define the label length accordingly.
+        self.expanded_inputs_length, self.targets_length = compute_input_and_target_lengths(
+            inputs_length=self.max_seq_length,
+            noise_density=self.noise_density,
+            mean_noise_span_length=self.mean_noise_span_length,
+        )
+
+    def train_dataloader(self):
+        from fengshen.data.universal_datamodule.universal_sampler import PretrainingSampler
+        from fengshen.data.universal_datamodule.universal_datamodule import get_consume_samples
+        # 采用自定义的sampler，确保继续训练能正确取到数据
+        consumed_samples = get_consume_samples(self)
+        batch_sampler = PretrainingSampler(
+            total_samples=len(self.train_dataset),
+            consumed_samples=consumed_samples,
+            micro_batch_size=self.hparams.train_batchsize,
+            data_parallel_rank=self.trainer.global_rank,
+            data_parallel_size=self.trainer.world_size,
+        )
+        return DataLoader(
+            self.train_dataset,
+            batch_sampler=batch_sampler,
+            pin_memory=True,
+            num_workers=self.hparams.dataloader_num_workers,
+            collate_fn=self.collate_fn,
+        )
+
+    def val_dataloader(self):
+        sampler = torch.utils.data.distributed.DistributedSampler(
+            self.test_dataset, shuffle=False)
+        return DataLoader(
+            self.test_dataset,
+            sampler=sampler,
+            shuffle=False,
+            batch_size=self.hparams.valid_batchsize,
+            pin_memory=True,
+            num_workers=self.hparams.dataloader_num_workers,
+            collate_fn=self.collate_fn,
+        )
+
+    def predict_dataloader(self):
+        sampler = torch.utils.data.distributed.DistributedSampler(
+            self.test_dataset, shuffle=False)
+        return DataLoader(
+            self.test_data,
+            sampler=sampler,
+            shuffle=False,
+            batch_size=self.hparams.valid_batchsize,
+            pin_memory=True,
+            num_workers=self.hparams.dataloader_num_workers,
+            collate_fn=self.collate_fn,
+        )
+
+    def collate_fn(self, examples):
+        # convert list to dict and tensorize input
+        batch = BatchEncoding(
+            {k: np.array([examples[i][k] for i in range(len(examples))])
+             for k, v in examples[0].items()}
+        )
+
+        input_ids = np.array(batch['input_ids'])
+        batch_size, expanded_input_length = input_ids.shape
+        mask_indices = np.asarray([self.random_spans_noise_mask(
+            expanded_input_length) for i in range(batch_size)])
+        labels_mask = ~mask_indices
+
+        input_ids_sentinel = self.create_sentinel_ids(
+            mask_indices.astype(np.int8))
+        labels_sentinel = self.create_sentinel_ids(labels_mask.astype(np.int8))
+
+        batch["input_ids"] = self.filter_input_ids(
+            input_ids, input_ids_sentinel)
+        batch["labels"] = self.filter_input_ids(input_ids, labels_sentinel)
+
+        if batch["input_ids"].shape[-1] != self.max_seq_length:
+            raise ValueError(
+                f"`input_ids` are incorrectly preprocessed. `input_ids` length is \
+                    {batch['input_ids'].shape[-1]}, but should be {self.targets_length}."
+            )
+
+        if batch["labels"].shape[-1] != self.targets_length:
+            raise ValueError(
+                f"`labels` are incorrectly preprocessed. `labels` length is \
+                    {batch['labels'].shape[-1]}, but should be {self.targets_length}."
+            )
+
+        batch["decoder_input_ids"] = self.shift_tokens_right(
+            batch["labels"], self.pad_token_id, self.decoder_start_token_id
+        )
+
+        for k, v in batch.items():
+            batch[k] = torch.tensor(v)
+            # print(k, batch[k], self.tokenizer.batch_decode(batch[k]), '\n', flush=True)
+        return batch
+
+    def create_sentinel_ids(self, mask_indices):
+        """
+        Sentinel ids creation given the indices that should be masked.
+        The start indices of each mask are replaced by the sentinel ids in increasing
+        order. Consecutive mask indices to be deleted are replaced with `-1`.
+        """
+        start_indices = mask_indices - \
+            np.roll(mask_indices, 1, axis=-1) * mask_indices
+        start_indices[:, 0] = mask_indices[:, 0]
+
+        sentinel_ids = np.where(start_indices != 0, np.cumsum(
+            start_indices, axis=-1), start_indices)
+        sentinel_ids = np.where(
+            sentinel_ids != 0, (self.vocab_size - sentinel_ids), 0)
+        sentinel_ids -= mask_indices - start_indices
+
+        return sentinel_ids
+
+    def filter_input_ids(self, input_ids, sentinel_ids):
+        """
+        Puts sentinel mask on `input_ids` and fuse consecutive mask tokens into a single mask token by deleting.
+        This will reduce the sequence length from `expanded_inputs_length` to `input_length`.
+        """
+        batch_size = input_ids.shape[0]
+
+        input_ids_full = np.where(sentinel_ids != 0, sentinel_ids, input_ids)
+        # input_ids tokens and sentinel tokens are >= 0, tokens < 0 are
+        # masked tokens coming after sentinel tokens and should be removed
+        input_ids = input_ids_full[input_ids_full >=
+                                   0].reshape((batch_size, -1))
+        input_ids = np.concatenate(
+            [input_ids, np.full((batch_size, 1), self.eos_token_id, dtype=np.int32)], axis=-1
+        )
+        return input_ids
+
+    # Copied from transformers.models.bart.modeling_flax_bart.shift_tokens_right
+    def shift_tokens_right(self, input_ids: np.array, pad_token_id: int, decoder_start_token_id: int) -> np.ndarray:
+        """
+        Shift input ids one token to the right.
+        """
+        shifted_input_ids = np.zeros_like(input_ids)
+        shifted_input_ids[:, 1:] = input_ids[:, :-1]
+        shifted_input_ids[:, 0] = decoder_start_token_id
+
+        shifted_input_ids = np.where(
+            shifted_input_ids == -100, pad_token_id, shifted_input_ids)
+        return shifted_input_ids
+
+    def random_spans_noise_mask(self, length):
+        """This function is copy of `random_spans_helper <https://github.com/google-research/text-to-text-transfer-transformer/
+        blob/84f8bcc14b5f2c03de51bd3587609ba8f6bbd1cd/t5/data/preprocessors.py#L2682>`__ .
+        Noise mask consisting of random spans of noise tokens.
+        The number of noise tokens and the number of noise spans and non-noise spans
+        are determined deterministically as follows:
+        num_noise_tokens = round(length * noise_density)
+        num_nonnoise_spans = num_noise_spans = round(num_noise_tokens / mean_noise_span_length)
+        Spans alternate between non-noise and noise, beginning with non-noise.
+        Subject to the above restrictions, all masks are equally likely.
+        Args:
+            length: an int32 scalar (length of the incoming token sequence)
+            noise_density: a float - approximate density of output mask
+            mean_noise_span_length: a number
+        Returns:
+            a boolean tensor with shape [length]
+        """
+
+        orig_length = length
+
+        num_noise_tokens = int(np.round(length * self.noise_density))
+        # avoid degeneracy by ensuring positive numbers of noise and nonnoise tokens.
+        num_noise_tokens = min(max(num_noise_tokens, 1), length - 1)
+        num_noise_spans = int(
+            np.round(num_noise_tokens / self.mean_noise_span_length))
+
+        # avoid degeneracy by ensuring positive number of noise spans
+        num_noise_spans = max(num_noise_spans, 1)
+        num_nonnoise_tokens = length - num_noise_tokens
+
+        # pick the lengths of the noise spans and the non-noise spans
+        def _random_segmentation(num_items, num_segments):
+            """Partition a sequence of items randomly into non-empty segments.
+            Args:
+                num_items: an integer scalar > 0
+                num_segments: an integer scalar in [1, num_items]
+            Returns:
+                a Tensor with shape [num_segments] containing positive integers that add
+                up to num_items
+            """
+            mask_indices = np.arange(num_items - 1) < (num_segments - 1)
+            np.random.shuffle(mask_indices)
+            first_in_segment = np.pad(mask_indices, [[1, 0]])
+            segment_id = np.cumsum(first_in_segment)
+            # count length of sub segments assuming that list is sorted
+            _, segment_length = np.unique(segment_id, return_counts=True)
+            return segment_length
+
+        noise_span_lengths = _random_segmentation(
+            num_noise_tokens, num_noise_spans)
+        nonnoise_span_lengths = _random_segmentation(
+            num_nonnoise_tokens, num_noise_spans)
+
+        interleaved_span_lengths = np.reshape(
+            np.stack([nonnoise_span_lengths, noise_span_lengths],
+                     axis=1), [num_noise_spans * 2]
+        )
+        span_starts = np.cumsum(interleaved_span_lengths)[:-1]
+        span_start_indicator = np.zeros((length,), dtype=np.int8)
+        span_start_indicator[span_starts] = True
+        span_num = np.cumsum(span_start_indicator)
+        is_noise = np.equal(span_num % 2, 1)
+
+        return is_noise[:orig_length]
+
+
+class TaskT5Dataset(Dataset):
+    def __init__(self, data_path, args):
+        super().__init__()
+        self.max_length = args.max_seq_length
+        if args.tokenizer_type == 't5_tokenizer':
+            self.tokenizer = MT5Tokenizer.from_pretrained(args.pretrained_model_path)
+        else:
+            self.tokenizer = BertTokenizer.from_pretrained(args.pretrained_model_path)
+        self.data = self.load_data(data_path)
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, index):
+        return self.encode(self.data[index])
+
+    def load_data(self, data_path):
+        samples = []
+        with open(data_path, 'r', encoding='utf8') as f:
+            lines = f.readlines()
+            for line in tqdm(lines):
+                samples.append(json.loads(line))
+        return samples
+
+    def encode(self, item):
+        if item["textb"] != "":
+            text = item['question'] + '，'.join(item['choice'])+'。' + f"""{item["texta"]}""" + f"""{item["textb"]}"""
+        else:
+            text = f"""{item["question"]}""" + "，".join(item["choice"]) + "。" + f"""{item["texta"]}"""
+        label = item['answer']
+        encode_dict = self.tokenizer.encode_plus(text, max_length=self.max_length, padding='max_length',
+                                                 truncation=True, return_tensors='pt')
+        decode_dict = self.tokenizer.encode_plus(label, max_length=16, padding='max_length',
+                                                 truncation=True)
+
+        answer_token = []
+        max_label_len = 0
+        choice_encode = []  # 用来确定模型生成的最大长度
+        for a in item['choice']:
+            answer_encode = self.tokenizer.encode(a)
+            choice_encode.append(answer_encode)
+            if len(answer_encode) > max_label_len:
+                max_label_len = len(answer_encode)
+            for an in answer_encode:
+                if an not in answer_token:
+                    answer_token.append(an)
+
+        # bad_words_ids = [[i] for i in range(self.tokenizer.vocab_size) if i not in answer_token] #不生成这些token
+
+        # while len(bad_words_ids)<self.tokenizer.vocab_size:
+        #     bad_words_ids.append(bad_words_ids[0])
+
+        # bad_words_ids = [[423],[67],[878]]
+
+        encode_sent = encode_dict['input_ids'].squeeze()
+        attention_mask = encode_dict['attention_mask'].squeeze()
+        target = decode_dict['input_ids']
+        labels = torch.tensor(target)
+        labels[target == self.tokenizer.pad_token_id] = -100
+
+        return {
+            "input_ids": torch.tensor(encode_sent).long(),
+            "attention_mask": torch.tensor(attention_mask).float(),
+            "labels": torch.tensor(target).long(),
+            "force_words_ids": answer_token,
+        }
+
+
+class TaskT5DataModel(pl.LightningDataModule):
+    @staticmethod
+    def add_data_specific_args(parent_args):
+        parser = parent_args.add_argument_group('TaskT5DataModel')
+        parser.add_argument('--dataset_num_workers', default=8, type=int)
+        parser.add_argument('--dataloader_num_workers', default=4, type=int)
+        parser.add_argument(
+            '--train_data_path', default='wudao_180g_mt5_tokenized', type=str)
+        parser.add_argument(
+            '--valid_data_path', default='wudao_180g_mt5_tokenized', type=str)
+        parser.add_argument('--train_batchsize', default=2, type=int)
+        parser.add_argument('--valid_batchsize', default=2, type=int)
+        parser.add_argument('--train_split_size', default=None, type=float)
+        parser.add_argument('--tokenizer_type', default='t5_tokenizer', choices=['t5_tokenizer', 'bert_tokenizer'])
+        parser.add_argument('--text_column_name', default='text')
+        parser.add_argument('--remove_columns', nargs='+', default=[])
+        return parent_args
+
+    def __init__(self, args):
+        super().__init__()
+        self.save_hyperparameters(args)
+        self.train_dataset = TaskT5Dataset(args.train_data_path, args)
+        self.valid_dataset = TaskT5Dataset(args.valid_data_path, args)
+
+    def train_dataloader(self):
+        from fengshen.data.universal_datamodule.universal_sampler import PretrainingSampler
+        from fengshen.data.universal_datamodule.universal_datamodule import get_consume_samples
+        # 采用自定��的sampler，确保继续训练能正确取到数据
+        consumed_samples = get_consume_samples(self)
+        # batch_sampler = PretrainingRandomSampler(
+        batch_sampler = PretrainingSampler(
+            total_samples=len(self.train_dataset),
+            consumed_samples=consumed_samples,
+            micro_batch_size=self.hparams.train_batchsize,
+            data_parallel_rank=self.trainer.global_rank,
+            data_parallel_size=self.trainer.world_size,
+        )
+        # epoch=self.trainer.current_epoch
+        # )
+        return DataLoader(
+            self.train_dataset,
+            batch_sampler=batch_sampler,
+            pin_memory=True,
+            num_workers=self.hparams.dataloader_num_workers
+        )
+
+    def val_dataloader(self):
+        sampler = torch.utils.data.distributed.DistributedSampler(
+            self.valid_dataset, shuffle=False)
+        return DataLoader(
+            self.valid_dataset,
+            sampler=sampler,
+            shuffle=False,
+            batch_size=self.hparams.valid_batchsize,
+            pin_memory=True,
+            num_workers=self.hparams.dataloader_num_workers
+        )

fengshen/data/t5_dataloader/t5_gen_datasets.py

@@ -0,0 +1,391 @@
+# -*- encoding: utf-8 -*-
+'''
+@File    :   t5_gen_datasets.py
+@Time    :   2022/10/24 19:29
+@Author  :   He Junqing
+@Version :   1.0
+@Contact :   hejunqing@idea.edu.cn
+@License :   (C)Copyright 2022-2023, CCNL-IDEA
+'''
+
+from logging import exception
+from transformers import (
+    BertTokenizer,
+    MT5Config,
+    MT5Tokenizer,
+    MT5ForConditionalGeneration,
+)
+import torch
+from torch.utils.data import Dataset, DataLoader
+from torch.nn.utils.rnn import pad_sequence
+import pytorch_lightning as pl
+import numpy as np
+import sys
+
+sys.path.append("../../")
+
+special_token_dict = {
+    "additional_special_tokens": [
+        "[CTSTART]",
+        "[CTEND]",
+        "[SEP]",
+        "[KNSTART]",
+        "[KNEND]",
+    ]
+}
+
+
+class DialogDataset(Dataset):
+    def __init__(self, data_path, args, data, load_data_type=1) -> None:
+        super().__init__()
+
+        if args.tokenizer_type == "t5_tokenizer":
+            self.tokenizer = MT5Tokenizer.from_pretrained(
+                args.pretrained_model_path)
+            if len(self.tokenizer) == 32596:
+                self.tokenizer.add_special_tokens(special_token_dict)
+                print(
+                    "add special tokens to tokenizer,vocab size:",
+                    len(self.tokenizer)
+                )
+                self.model = MT5ForConditionalGeneration.from_pretrained(
+                    args.pretrained_model_path
+                )
+                self.model.resize_token_embeddings(len(self.tokenizer))
+                self.model.save_pretrained(args.new_vocab_path)
+                self.tokenizer.save_pretrained(
+                    args.new_vocab_path)
+        else:
+            self.tokenizer = BertTokenizer.from_pretrained(
+                args.pretrained_model_path)
+
+        self.load_data_type = load_data_type
+        self.data_split = data
+        self.num_workers = args.preprocessing_num_workers
+        self.max_seq_length = args.max_seq_length
+        self.max_knowledge_length = args.max_knowledge_length
+        self.max_target_length = args.max_target_length
+
+        # tokenizer config
+        self.config = MT5Config.from_pretrained(args.pretrained_model_path)
+        self.decoder_start_token_id = self.config.decoder_start_token_id
+        self.eos_token_id = self.config.eos_token_id
+        self.vocab_size = self.config.vocab_size
+        # print(self.tokenizer.decode([2]))
+
+        # load from raw data or hf dataset
+
+        if self.load_data_type == 0:
+            self.data = self.load_data(data_path)
+        elif self.load_data_type == 1:
+            self.data = self.load_packed_data(data_path)
+        else:  # for testing
+            self.data = data_path
+
+    def load_packed_data(self, data_path):
+        from fengshen.data.fs_datasets import load_dataset
+
+        samples = load_dataset(data_path,
+                               num_proc=self.num_workers)[self.data_split]
+        tokenized_samples = samples.map(
+            self.regular_tokenize, batched=False,
+            num_proc=self.num_workers
+        )
+
+        return tokenized_samples
+
+    def load_data(self, data_path):
+        """
+        load data from raw data
+        return untokoenized data
+        """
+        from datasets import load_dataset
+
+        ds = load_dataset("json", data_files=data_path)['train']
+        samples = ds.map(self.regular_tokenize, batched=False, num_proc=self.num_workers
+                         )
+        return samples
+
+    def __getitem__(self, index):
+        return self.data[index]
+
+    def __len__(self):
+        return len(self.data)
+
+    def regular_tokenize(self, sample):
+        # print(len(sample['context']))
+        context_ids = self.tokenizer(
+            sample["context"],
+            add_special_tokens=True,
+            return_attention_mask=False,
+            return_token_type_ids=True,
+        )
+
+        context_types = self.get_token_type(
+            sample["context"], context_ids["token_type_ids"]
+        )
+        # print('context',sample['context'])
+        # print('context_ids',context_ids['input_ids'])
+        knowledge_ids = self.tokenizer.encode(
+            sample["knowledge"], add_special_tokens=False
+        )
+        # print('knowledge_ids',knowledge_ids)
+        if isinstance(knowledge_ids, int):
+            knowledge_ids = [knowledge_ids]
+        target_ids = self.tokenizer.encode(
+            sample["target"],
+            add_special_tokens=False,
+            max_length=self.max_target_length - 1,
+            truncation=True,
+        )
+        # print('target',sample['target'])
+        # print('target_ids',target_ids)
+        # print('decode target',self.tokenizer.decode(target_ids))
+        # truncate
+
+        knowledge_ids = (
+            [self.tokenizer.convert_tokens_to_ids("[KNSTART]")]
+            + knowledge_ids[: self.max_knowledge_length - 2]
+            + [self.tokenizer.convert_tokens_to_ids("[KNEND]")]
+        )
+        l_kn = len(knowledge_ids)
+        knowledge_types = [2] * l_kn
+
+        flatten_context = []
+        for line in context_ids["input_ids"]:
+            flatten_context.extend(line)
+        l_ct = min(len(flatten_context), self.max_seq_length - l_kn - 2)
+        context_ids = (
+            [self.tokenizer.convert_tokens_to_ids("[CTSTART]")]
+            + flatten_context[-l_ct:]
+            + [self.tokenizer.convert_tokens_to_ids("[CTEND]")]
+        )
+
+        context_types = context_types[-l_ct:] + [0]
+        context_types.insert(0, context_types[0])
+        assert len(context_ids) == len(
+            context_types
+        ), "len of context ids and token types unmatch, context:{},ids:{} types:{},len {}:{}".format(
+            sample["context"],
+            context_ids,
+            context_types,
+            len(context_ids),
+            len(context_types),
+        )
+
+        try:
+            target_ids = target_ids + [self.eos_token_id]
+        except exception:
+            print(sample["target"], target_ids, self.eos_token_id)
+
+        tokenized = {}
+        tokenized["input_ids"] = np.array(context_ids + knowledge_ids, dtype=np.int32)
+        tokenized["token_types"] = np.array(
+            context_types + knowledge_types, dtype=np.int32
+        )
+        tokenized["attention_mask"] = np.ones(
+            len(context_types + knowledge_types), dtype=np.int8
+        )
+        tokenized["labels"] = np.array(target_ids, dtype=np.int32)
+
+        return tokenized
+
+    def get_token_type(self, context, tokentypes=None):
+        # token_type fail in tokenizer, all zero
+        context_token_types = []
+        for i, line in enumerate(context):
+            if tokentypes:
+                if i % 2 == 0:
+                    token_type = [0] * len(tokentypes[i])
+                else:
+                    token_type = [1] * len(tokentypes[i])
+            else:
+                if i % 2 == 0:
+                    token_type = [0] * (1 + len(line))
+                else:
+                    token_type = [1] * (1 + len(line))
+
+            context_token_types.extend(token_type)
+
+        return context_token_types
+
+
+class DialogDataModel(pl.LightningDataModule):
+    @staticmethod
+    def add_data_specific_args(parent_args):
+        parser = parent_args.add_argument_group("SuperviseT5DataModel")
+        parser.add_argument("--dataset_num_workers", default=8, type=int)
+        parser.add_argument("--dataloader_num_workers", default=4, type=int)
+        parser.add_argument("--train_data_path", default="dialog_4g_test", type=str)
+        parser.add_argument(
+            "--valid_data_path", default="wudao_180g_mt5_tokenized", type=str
+        )
+        parser.add_argument("--train_batchsize", default=2, type=int)
+        parser.add_argument("--valid_batchsize", default=2, type=int)
+        parser.add_argument("--max_seq_length", default=512, type=int)
+        parser.add_argument("--max_knowledge_length", default=128, type=int)
+        parser.add_argument("--max_target_length", default=128, type=int)
+
+        return parent_args
+
+    def __init__(self, args):
+        super().__init__()
+        self.save_hyperparameters(args)
+        self.load_data(args)
+        self.epochs = args.max_epochs
+
+    def load_data(self, args):
+        if args.train_split_size is not None:
+            from fengshen.data.fs_datasets import load_dataset
+
+            data_splits = load_dataset(
+                args.train_data_path, num_proc=args.dataset_num_workers
+            )
+            train_split = data_splits['train']
+            test_split = data_splits['test']
+            print('train:', train_split, '\ntest_data:', test_split)
+            self.train_dataset = DialogDataset(
+                args.train_data_path, args, load_data_type=1, data="train"
+            )
+            self.test_dataset = DialogDataset(
+                args.train_data_path, args, load_data_type=1, data="test"
+            )
+        else:
+            self.train_data = DialogDataset(
+                args.train_data_path, args, load_data_type=1
+            )
+
+        self.config = MT5Config.from_pretrained(args.pretrained_model_path)
+        self.pad_token_id = self.config.pad_token_id
+        self.decoder_start_token_id = self.config.decoder_start_token_id
+        print("bos id:", self.decoder_start_token_id)
+
+    def collate_fn(self, samples):
+        batch = {
+            k: [
+                torch.tensor(samples[i][k], dtype=torch.int64)
+                for i in range(len(samples))
+            ]
+            for k in ["input_ids", "token_types", "attention_mask", "labels"]
+        }
+
+        # print(batch)
+        for k, v in batch.items():
+            if k != "labels":
+                batch[k] = pad_sequence(
+                    v, batch_first=True, padding_value=self.pad_token_id
+                )
+            else:
+                batch[k] = pad_sequence(v, batch_first=True, padding_value=-100)
+        batch["decoder_input_ids"] = torch.tensor(
+            self.shift_tokens_right(
+                batch["labels"], self.pad_token_id, self.decoder_start_token_id
+            ),
+            dtype=torch.long,
+        )
+        return batch
+
+    def shift_tokens_right(
+        self, input_ids: np.array, pad_token_id: int, decoder_start_token_id: int
+    ) -> np.ndarray:
+        """
+        Shift input ids one token to the right.
+        """
+        shifted_input_ids = np.zeros_like(input_ids)
+        shifted_input_ids[:, 1:] = input_ids[:, :-1]
+        shifted_input_ids[:, 0] = decoder_start_token_id
+
+        shifted_input_ids = np.where(
+            shifted_input_ids == -100, pad_token_id, shifted_input_ids
+        )
+        return shifted_input_ids
+
+    def train_dataloader(self):
+        from fengshen.data.universal_datamodule.universal_sampler import (
+            PretrainingRandomSampler,
+        )
+        from fengshen.data.universal_datamodule.universal_datamodule import (
+            get_consume_samples,
+        )
+
+        # 采用自定义的sampler，确保继续训练能正确取到数据
+        consumed_samples = get_consume_samples(self)
+        batch_sampler = PretrainingRandomSampler(
+            epoch=self.epochs,
+            total_samples=len(self.train_dataset),
+            consumed_samples=consumed_samples,
+            micro_batch_size=self.hparams.train_batchsize,
+            data_parallel_rank=self.trainer.global_rank,  # gpu idx
+            data_parallel_size=self.trainer.world_size,  # gpu num
+        )
+        return DataLoader(
+            self.train_dataset,
+            batch_sampler=batch_sampler,
+            pin_memory=True,
+            num_workers=self.hparams.dataloader_num_workers,
+            collate_fn=self.collate_fn,
+        )
+
+    def val_dataloader(self):
+        sampler = torch.utils.data.distributed.DistributedSampler(
+            self.test_dataset, shuffle=False
+        )
+        return DataLoader(
+            self.test_dataset,
+            sampler=sampler,
+            shuffle=False,
+            batch_size=self.hparams.valid_batchsize,
+            pin_memory=True,
+            num_workers=self.hparams.dataloader_num_workers,
+            collate_fn=self.collate_fn,
+        )
+
+    def predict_dataloader(self):
+        sampler = torch.utils.data.distributed.DistributedSampler(
+            self.test_dataset, shuffle=False
+        )
+        return DataLoader(
+            self.test_dataset,
+            sampler=sampler,
+            shuffle=False,
+            batch_size=self.hparams.valid_batchsize,
+            pin_memory=True,
+            num_workers=self.hparams.dataloader_num_workers,
+            collate_fn=self.collate_fn,
+        )
+
+
+if __name__ == "__main__":
+    # test
+    import argparse
+
+    total_parser = argparse.ArgumentParser("DATASET parser")
+    total_parser.add_argument(
+        "--tokenizer_type",
+        default="t5_tokenizer",
+        choices=["bert_tokenizer", "t5_tokenizer"],
+    )
+    total_parser.add_argument("--preprocessing_num_workers", default="10", type=int)
+    total_parser.add_argument(
+        "--new_vocab_path",
+        default="/cognitive_comp/hejunqing/projects/Dialog_pretrain/randeng_t5_newvocab_784M",
+        type=str,
+    )
+    total_parser.add_argument("--train_split_size", default=0.995, type=int)
+    total_parser.add_argument(
+        "--pretrained_model_path",
+        default="/cognitive_comp/hejunqing/projects/Dialog_pretrain/randeng_t5_newvocab_784M",
+    )
+    total_parser = DialogDataModel.add_data_specific_args(total_parser)
+    args = total_parser.parse_args()
+    dl = DialogDataModel(args)
+
+    for i in range(5):
+        for batch in dl.train_dataloader():
+            print(batch)
+            print(batch["input_ids"])
+            print(batch["token_types"])
+            print(batch["decoder_input_ids"])
+            print(batch["labels"])
+
+    print("test finish")

fengshen/data/taiyi_stable_diffusion_datasets/taiyi_datasets.py

@@ -0,0 +1,173 @@
+from torch.utils.data import Dataset, ConcatDataset
+import os
+from concurrent.futures import ProcessPoolExecutor
+import pandas as pd
+
+
+def add_data_args(parent_args):
+    parser = parent_args.add_argument_group('taiyi stable diffusion data args')
+    # 支持传入多个路径，分别加载
+    parser.add_argument(
+        "--datasets_path", type=str, default=None, required=True, nargs='+',
+        help="A folder containing the training data of instance images.",
+    )
+    parser.add_argument(
+        "--datasets_type", type=str, default=None, required=True, choices=['txt', 'csv', 'fs_datasets'], nargs='+',
+        help="dataset type, txt or csv, same len as datasets_path",
+    )
+    parser.add_argument(
+        "--resolution", type=int, default=512,
+        help=(
+            "The resolution for input images, all the images in the train/validation dataset will be resized to this"
+            " resolution"
+        ),
+    )
+    parser.add_argument(
+        "--center_crop", action="store_true", default=False,
+        help="Whether to center crop images before resizing to resolution"
+    )
+    parser.add_argument("--thres", type=float, default=0.2)
+    return parent_args
+
+
+class TXTDataset(Dataset):
+    # 添加Txt数据集读取，主要是针对Zero23m数据集。
+    def __init__(self,
+                 foloder_name,
+                 thres=0.2):
+        super().__init__()
+        # print(f'Loading folder data from {foloder_name}.')
+        self.image_paths = []
+        '''
+        暂时没有开源这部分文件
+        score_data = pd.read_csv(os.path.join(foloder_name, 'score.csv'))
+        img_path2score = {score_data['image_path'][i]: score_data['score'][i]
+                          for i in range(len(score_data))}
+        '''
+        # print(img_path2score)
+        # 这里都存的是地址，避免初始化时间过多。
+        for each_file in os.listdir(foloder_name):
+            if each_file.endswith('.jpg'):
+                self.image_paths.append(os.path.join(foloder_name, each_file))
+
+        # print('Done loading data. Len of images:', len(self.image_paths))
+
+    def __len__(self):
+        return len(self.image_paths)
+
+    def __getitem__(self, idx):
+        img_path = str(self.image_paths[idx])
+        caption_path = img_path.replace('.jpg', '.txt')  # 图片名称和文本名称一致。
+        with open(caption_path, 'r') as f:
+            caption = f.read()
+        return {'img_path': img_path, 'caption': caption}
+
+
+# NOTE 加速读取数据，直接用原版的，在外部使用并行读取策略。30min->3min
+class CSVDataset(Dataset):
+    def __init__(self,
+                 input_filename,
+                 image_root,
+                 img_key,
+                 caption_key,
+                 thres=0.2):
+        super().__init__()
+        # logging.debug(f'Loading csv data from {input_filename}.')
+        print(f'Loading csv data from {input_filename}.')
+        self.images = []
+        self.captions = []
+
+        if input_filename.endswith('.csv'):
+            # print(f"Load Data from{input_filename}")
+            df = pd.read_csv(input_filename, index_col=0, on_bad_lines='skip')
+            print(f'file {input_filename} datalen {len(df)}')
+            # 这个图片的路径也需要根据数据集的结构稍微做点修改
+            self.images.extend(df[img_key].tolist())
+            self.captions.extend(df[caption_key].tolist())
+        self.image_root = image_root
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, idx):
+        img_path = os.path.join(self.image_root, str(self.images[idx]))
+        return {'img_path': img_path, 'caption': self.captions[idx]}
+
+
+def if_final_dir(path: str) -> bool:
+    # 如果当前目录有一个文件，那就算是终极目录
+    for f in os.scandir(path):
+        if f.is_file():
+            return True
+    return False
+
+
+def process_pool_read_txt_dataset(args,
+                                  input_root=None,
+                                  thres=0.2):
+    p = ProcessPoolExecutor(max_workers=20)
+    all_datasets = []
+    res = []
+
+    # 遍历该目录下所有的子目录
+    def traversal_files(path: str):
+        list_subfolders_with_paths = [f.path for f in os.scandir(path) if f.is_dir()]
+        for dir_path in list_subfolders_with_paths:
+            if if_final_dir(dir_path):
+                res.append(p.submit(TXTDataset,
+                                    dir_path,
+                                    thres))
+            else:
+                traversal_files(dir_path)
+    traversal_files(input_root)
+    p.shutdown()
+    for future in res:
+        all_datasets.append(future.result())
+    dataset = ConcatDataset(all_datasets)
+    return dataset
+
+
+def process_pool_read_csv_dataset(args,
+                                  input_root,
+                                  thres=0.20):
+    # here input_filename is a directory containing a CSV file
+    all_csvs = os.listdir(os.path.join(input_root, 'release'))
+    image_root = os.path.join(input_root, 'images')
+    # csv_with_score = [each for each in all_csvs if 'score' in each]
+    all_datasets = []
+    res = []
+    p = ProcessPoolExecutor(max_workers=150)
+    for path in all_csvs:
+        each_csv_path = os.path.join(input_root, 'release', path)
+        res.append(p.submit(CSVDataset,
+                            each_csv_path,
+                            image_root,
+                            img_key="name",
+                            caption_key="caption",
+                            thres=thres))
+    p.shutdown()
+    for future in res:
+        all_datasets.append(future.result())
+    dataset = ConcatDataset(all_datasets)
+    return dataset
+
+
+def load_data(args, global_rank=0):
+    assert len(args.datasets_path) == len(args.datasets_type), \
+        "datasets_path num not equal to datasets_type"
+    all_datasets = []
+    for path, type in zip(args.datasets_path, args.datasets_type):
+        if type == 'txt':
+            all_datasets.append(process_pool_read_txt_dataset(
+                args, input_root=path, thres=args.thres))
+        elif type == 'csv':
+            all_datasets.append(process_pool_read_csv_dataset(
+                args, input_root=path, thres=args.thres))
+        elif type == 'fs_datasets':
+            from fengshen.data.fs_datasets import load_dataset
+            all_datasets.append(load_dataset(path, num_proc=args.num_workers,
+                                             thres=args.thres, global_rank=global_rank)['train'])
+        else:
+            raise ValueError('unsupport dataset type: %s' % type)
+        print(f'load datasset {type} {path} len {len(all_datasets[-1])}')
+    return {'train': ConcatDataset(all_datasets)}

fengshen/data/task_dataloader/__init__.py

@@ -0,0 +1,3 @@
+# coding=utf-8
+from .task_datasets import LCSTSDataModel, LCSTSDataset
+__all__ = ['LCSTSDataModel', 'LCSTSDataset']

fengshen/data/task_dataloader/medicalQADataset.py

@@ -0,0 +1,137 @@
+# coding=utf8
+import os
+import pytorch_lightning as pl
+from torch.utils.data import DataLoader, Dataset
+from tqdm import tqdm
+from transformers import AutoTokenizer
+
+
+class GPT2QADataset(Dataset):
+    '''
+    Dataset Used for yuyuan medical qa task.
+    Just surpport small datasets, when deal with large datasets it may be slowly.
+    for large datasets please use mmapdatasets(doing)
+    '''
+
+    def __init__(self, data_path, name, args):
+        super().__init__()
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            args.pretrained_model_path)
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.add_special_tokens({'pad_token': '<|endoftext|>'})
+        self.data_size = os.path.getsize(data_path)/1024/1024/1024
+        self.data_type_name = name
+        self.data = self.load_data(data_path)
+        self.max_seq_length = args.max_seq_length
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, index):
+        return self.encode(self.data[index])
+
+    def load_data(self, data_path):
+        # 有进度条展示
+        if self.data_size <= 5:
+            with open(data_path, "rt", encoding='utf8') as f:
+                lines = f.readlines()
+            total_num = len(lines)
+            data_gen = lines
+        else:
+            data_gen = open(data_path, "rt", encoding='utf8')
+            total_num = None
+
+        data = []
+        with tqdm(total=total_num, desc=f'{self.data_type_name}处理进度', mininterval=0.3) as bar:
+            for idx, line in enumerate(data_gen):
+                data.append(self.data_parse(line))
+                bar.update()
+
+        if self.data_size > 5:
+            data_gen.close()
+        return data
+
+    def data_parse(self, line):
+        """
+        解析不同格式的数据
+        """
+        dic = eval(line.strip())
+        return dic
+
+    def encode(self, item):
+        """
+        将数据转换成模型训练的输入
+        """
+        inputs_dict = self.tokenizer.encode_plus(item['Question']+item['answer'],
+                                                 max_length=self.max_seq_length, padding='max_length',
+                                                 truncation=True, return_tensors='pt')
+        target = inputs_dict['input_ids']
+        labels = target.clone().detach()
+        labels[target == self.tokenizer.pad_token_id] = -100
+        return {
+            "input_ids": inputs_dict['input_ids'].squeeze(),
+            "attention_mask": inputs_dict['attention_mask'].squeeze(),
+            "labels": labels.squeeze(),
+            "question": item['Question'],
+            "answer": item['answer']
+        }
+
+
+class GPT2QADataModel(pl.LightningDataModule):
+    @staticmethod
+    def add_data_specific_args(parent_args):
+        parser = parent_args.add_argument_group('GPT2QADataModel')
+        parser.add_argument('--data_dir', type=str, required=True)
+        parser.add_argument('--num_workers', default=2, type=int)
+        parser.add_argument('--train_data', default='train.txt', type=str)
+        parser.add_argument('--valid_data', default='valid.txt', type=str)
+        parser.add_argument('--test_data', default='test.txt', type=str)
+        parser.add_argument('--train_batchsize', type=int, required=True)
+        parser.add_argument('--valid_batchsize', type=int, required=True)
+        parser.add_argument('--max_seq_length', default=1024, type=int)
+        return parent_args
+
+    def __init__(self, args):
+        super().__init__()
+        self.args = args
+        self.train_batchsize = args.train_batchsize
+        self.valid_batchsize = args.valid_batchsize
+        if not args.do_eval_only:
+            self.train_data = GPT2QADataset(os.path.join(
+                args.data_dir, args.train_data), '训练集', args)
+            self.valid_data = GPT2QADataset(os.path.join(
+                args.data_dir, args.valid_data), '验证集', args)
+        self.test_data = GPT2QADataset(os.path.join(
+            args.data_dir, args.test_data), '测试集', args)
+
+    def train_dataloader(self):
+        return DataLoader(
+            self.train_data, shuffle=True,
+            batch_size=self.train_batchsize,
+            pin_memory=False, num_workers=self.args.num_workers)
+
+    def val_dataloader(self):
+        return DataLoader(self.valid_data, shuffle=False,
+                          batch_size=self.valid_batchsize,
+                          pin_memory=False, num_workers=self.args.num_workers)
+
+    def predict_dataloader(self):
+        return DataLoader(self.test_data, shuffle=False,
+                          batch_size=self.valid_batchsize, pin_memory=False,
+                          num_workers=self.args.num_workers)
+
+
+if __name__ == '__main__':
+    import argparse
+    modelfile = '/cognitive_comp/wuziwei/pretrained_model_hf/medical_v2'
+    datafile = '/cognitive_comp/wuziwei/task-data/medical_qa/medical_qa_train.txt'
+    parser = argparse.ArgumentParser(description='hf test', allow_abbrev=False)
+    group = parser.add_argument_group(title='test args')
+    group.add_argument('--pretrained-model-path', type=str, default=modelfile,
+                       help='Number of transformer layers.')
+    group.add_argument('--max-seq-length', type=int, default=1024)
+    args = parser.parse_args()
+
+    testml = GPT2QADataset(datafile, 'medical_qa', args=args)
+
+    print(testml[10])

fengshen/data/task_dataloader/task_datasets.py

@@ -0,0 +1,206 @@
+# coding=utf8
+from torch.utils.data import Dataset, DataLoader
+from tqdm import tqdm
+from transformers import AutoTokenizer
+import json
+import torch
+import pytorch_lightning as pl
+import os
+
+
+class AbstractCollator:
+    """
+    collector for summary task
+    """
+
+    def __init__(self, tokenizer, max_enc_length, max_dec_length, prompt):
+        self.tokenizer = tokenizer
+        self.max_enc_length = max_enc_length
+        self.max_dec_length = max_dec_length
+        self.prompt = prompt
+
+    def __call__(self, samples):
+
+        labels = []
+        attn_mask = []
+        # decoder_attn_mask = []
+        source_inputs = []
+        for sample in samples:
+            encode_dict = self.tokenizer.encode_plus(
+                self.prompt + sample['text'],
+                max_length=self.max_enc_length,
+                padding='max_length',
+                truncation=True,
+                return_tensors='pt')
+            decode_dict = self.tokenizer.encode_plus(
+                sample['summary'],
+                max_length=self.max_dec_length,
+                padding='max_length',
+                truncation=True,
+                return_tensors='pt')
+            source_inputs.append(encode_dict['input_ids'].squeeze())
+            labels.append(decode_dict['input_ids'].squeeze())
+            attn_mask.append(encode_dict['attention_mask'].squeeze())
+            # decoder_attn_mask.append(decode_dict['attention_mask'].squeeze())
+        # labels = torch.tensor(decode_dict['input'])
+
+        source_inputs = torch.stack(source_inputs)
+        labels = torch.stack(labels)
+        attn_mask = torch.stack(attn_mask)
+        # decoder_attn_mask = torch.stack(decoder_attn_mask)
+        # decode_input_idxs = shift_tokens_right(labels, self.tokenizer.pad_token_id, self.tokenizer.pad_token_id)
+        end_token_index = torch.where(labels == self.tokenizer.eos_token_id)[1]
+        for idx, end_idx in enumerate(end_token_index):
+            labels[idx][end_idx + 1:] = -100
+
+        return {
+            "input_ids": source_inputs,
+            "attention_mask": attn_mask,
+            "labels": labels,
+            "text": [sample['text'] for sample in samples],
+            "summary": [sample['summary'] for sample in samples]
+        }
+
+
+class LCSTSDataset(Dataset):
+    '''
+    Dataset Used for LCSTS summary task.
+    '''
+
+    def __init__(self, data_path, args):
+        super().__init__()
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            args.pretrained_model_path, use_fast=False)
+        self.data = self.load_data(data_path)
+        self.prompt = args.prompt
+        self.max_enc_length = args.max_enc_length
+        self.max_dec_length = args.max_dec_length
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, index):
+        return self.encode(self.data[index])
+
+    def load_data(self, data_path):
+        with open(data_path, "r", encoding='utf8') as f:
+            lines = f.readlines()
+        samples = []
+        for line in tqdm(lines):
+            obj = json.loads(line)
+            source = obj['text']
+            target = obj['summary']
+            samples.append({
+                "text": source,
+                "summary": target
+            })
+        return samples
+
+    def cal_data(self, data_path):
+        with open(data_path, "r", encoding='utf8') as f:
+            lines = f.readlines()
+        samples = []
+        enc_sizes = []
+        dec_sizes = []
+        for line in tqdm(lines):
+            obj = json.loads(line.strip())
+            source = obj['text']
+            target = obj['summary']
+            enc_input_ids = self.tokenizer.encode(source)
+            target = self.tokenizer.encode(target)
+            enc_sizes.append(len(enc_input_ids))
+            dec_sizes.append(len(target)-1)
+            samples.append({
+                "enc_input_ids": enc_input_ids,
+                "dec_input_ids": target[:-1],
+                "label_ids": target[1:]
+            })
+        max_enc_len = max(enc_sizes)
+        max_dec_len = max(dec_sizes)
+        import numpy as np
+        # mean of len(enc_input_ids): 74.68041911345998
+        # mean of len(dec_input_ids): 14.02265483791283
+        # max of len(enc_input_ids): 132
+        # max of len(dec_input_ids): 31
+        print('mean of len(enc_input_ids):', np.mean(enc_sizes),
+              'mean of len(dec_input_ids):', np.mean(dec_sizes),
+              'max of len(enc_input_ids):', max_enc_len,
+              'max of len(dec_input_ids):', max_dec_len)
+        return samples
+
+    def encode(self, item):
+        encode_dict = self.tokenizer.encode_plus(
+            self.prompt + item['text'],
+            max_length=self.max_enc_length,
+            padding='max_length',
+            truncation=True,
+            return_tensors='pt')
+        decode_dict = self.tokenizer.encode_plus(
+            item['summary'],
+            max_length=self.max_dec_length,
+            padding='max_length',
+            truncation=True)
+
+        target = decode_dict['input_ids']
+        # print('encode_dict shape:', encode_dict['input_ids'].shape)
+        labels = torch.tensor(target)
+        labels[target == self.tokenizer.pad_token_id] = -100
+        return {
+            "input_ids": encode_dict['input_ids'].squeeze(),
+            "attention_mask": encode_dict['attention_mask'].squeeze(),
+            "labels": labels.squeeze(),
+            "text": item['text'],
+            "summary": item['summary']
+        }
+
+
+class LCSTSDataModel(pl.LightningDataModule):
+    @staticmethod
+    def add_data_specific_args(parent_args):
+        parser = parent_args.add_argument_group('LCSTSDataModel')
+        parser.add_argument(
+            '--data_dir', default='/cognitive_comp/ganruyi/data_datasets_LCSTS_LCSTS/', type=str)
+        parser.add_argument('--num_workers', default=8, type=int)
+        parser.add_argument('--train_data', default='train.jsonl', type=str)
+        parser.add_argument('--valid_data', default='valid.jsonl', type=str)
+        parser.add_argument('--test_data', default='test_public.jsonl', type=str)
+        parser.add_argument('--train_batchsize', default=128, type=int)
+        parser.add_argument('--valid_batchsize', default=128, type=int)
+        parser.add_argument('--max_enc_length', default=128, type=int)
+        parser.add_argument('--max_dec_length', default=30, type=int)
+        parser.add_argument('--prompt', default='summarize:', type=str)
+        return parent_args
+
+    def __init__(self, args):
+        super().__init__()
+        self.args = args
+        self.train_batchsize = args.train_batchsize
+        self.valid_batchsize = args.valid_batchsize
+        if not args.do_eval_only:
+            self.train_data = LCSTSDataset(os.path.join(
+                args.data_dir, args.train_data), args)
+        self.valid_data = LCSTSDataset(os.path.join(
+            args.data_dir, args.valid_data), args)
+        self.test_data = LCSTSDataset(os.path.join(
+            args.data_dir, args.test_data), args)
+
+    def train_dataloader(self):
+        return DataLoader(self.train_data,
+                          shuffle=True,
+                          batch_size=self.train_batchsize,
+                          pin_memory=False,
+                          num_workers=self.args.num_workers)
+
+    def val_dataloader(self):
+        return DataLoader(self.valid_data,
+                          shuffle=False,
+                          batch_size=self.valid_batchsize,
+                          pin_memory=False,
+                          num_workers=self.args.num_workers)
+
+    def predict_dataloader(self):
+        return DataLoader(self.test_data,
+                          shuffle=False,
+                          batch_size=self.valid_batchsize,
+                          pin_memory=False,
+                          num_workers=self.args.num_workers)

fengshen/data/universal_datamodule/__init__.py

@@ -0,0 +1,4 @@
+from .universal_datamodule import UniversalDataModule
+from .universal_sampler import PretrainingSampler, PretrainingRandomSampler
+
+__all__ = ['UniversalDataModule', 'PretrainingSampler', 'PretrainingRandomSampler']

fengshen/data/universal_datamodule/universal_datamodule.py

@@ -0,0 +1,165 @@
+from pytorch_lightning import LightningDataModule
+from typing import Optional
+
+from torch.utils.data import DataLoader, DistributedSampler
+
+
+def get_consume_samples(data_model: LightningDataModule) -> int:
+    if hasattr(data_model.trainer.lightning_module, 'consumed_samples'):
+        consumed_samples = data_model.trainer.lightning_module.consumed_samples
+        print('get consumed samples from model: {}'.format(consumed_samples))
+    else:
+        world_size = data_model.trainer.world_size
+        consumed_samples = max(0, data_model.trainer.global_step - 1) * \
+            data_model.hparams.train_batchsize * world_size * data_model.trainer.accumulate_grad_batches
+        print('calculate consumed samples: {}'.format(consumed_samples))
+    return consumed_samples
+
+
+class UniversalDataModule(LightningDataModule):
+    @ staticmethod
+    def add_data_specific_args(parent_args):
+        parser = parent_args.add_argument_group('Universal DataModule')
+        parser.add_argument('--num_workers', default=8, type=int)
+        parser.add_argument('--dataloader_workers', default=2, type=int)
+        parser.add_argument('--train_batchsize', default=16, type=int)
+        parser.add_argument('--val_batchsize', default=16, type=int)
+        parser.add_argument('--test_batchsize', default=16, type=int)
+        parser.add_argument('--datasets_name', type=str, default=None)
+        parser.add_argument('--train_datasets_field', type=str, default='train')
+        parser.add_argument('--val_datasets_field', type=str, default='validation')
+        parser.add_argument('--test_datasets_field', type=str, default='test')
+        parser.add_argument('--train_file', type=str, default=None)
+        parser.add_argument('--val_file', type=str, default=None)
+        parser.add_argument('--test_file', type=str, default=None)
+        parser.add_argument('--raw_file_type', type=str, default='json')
+        parser.add_argument('--sampler_type', type=str,
+                            choices=['single',
+                                     'random'],
+                            default='random')
+        return parent_args
+
+    def __init__(
+        self,
+        tokenizer,
+        collate_fn,
+        args,
+        datasets=None,
+        **kwargs,
+    ):
+        super().__init__()
+        # 如果不传入datasets的名字，则可以在对象外部替换内部的datasets为模型需要的
+        if datasets is not None:
+            self.datasets = datasets
+        elif args.datasets_name is not None:
+            from fengshen.data.fs_datasets import load_dataset
+            print('---------begin to load datasets {}'.format(args.datasets_name))
+            self.datasets = load_dataset(
+                args.datasets_name, num_proc=args.num_workers)
+            print('---------ending load datasets {}'.format(args.datasets_name))
+        else:
+            print('---------begin to load datasets from local file')
+            from datasets import load_dataset
+            self.datasets = load_dataset(args.raw_file_type,
+                                         data_files={
+                                             args.train_datasets_field: args.train_file,
+                                             args.val_datasets_field: args.val_file,
+                                             args.test_datasets_field: args.test_file})
+            print('---------end to load datasets from local file')
+
+        self.tokenizer = tokenizer
+        self.collate_fn = collate_fn
+        self.save_hyperparameters(args)
+
+    def get_custom_sampler(self, ds):
+        from .universal_sampler import PretrainingRandomSampler
+        from .universal_sampler import PretrainingSampler
+        world_size = self.trainer.world_size
+        consumed_samples = get_consume_samples(self)
+        # use the user default sampler
+        if self.hparams.sampler_type == 'random':
+            return PretrainingRandomSampler(
+                total_samples=len(ds),
+                # consumed_samples cal by global steps
+                consumed_samples=consumed_samples,
+                micro_batch_size=self.hparams.train_batchsize,
+                data_parallel_rank=self.trainer.global_rank,
+                data_parallel_size=world_size,
+                epoch=self.trainer.current_epoch,
+            )
+        elif self.hparams.sampler_type == 'single':
+            return PretrainingSampler(
+                total_samples=len(ds),
+                # consumed_samples cal by global steps
+                consumed_samples=consumed_samples,
+                micro_batch_size=self.hparams.train_batchsize,
+                data_parallel_rank=self.trainer.global_rank,
+                data_parallel_size=world_size,
+            )
+        else:
+            raise Exception('Unknown sampler type: {}'.format(self.hparams.sampler_type))
+
+    def setup(self, stage: Optional[str] = None) -> None:
+        return
+
+    def train_dataloader(self):
+        ds = self.datasets[self.hparams.train_datasets_field]
+
+        collate_fn = self.collate_fn
+        if hasattr(ds, 'collate_fn'):
+            collate_fn = ds.collate_fn
+
+        if self.hparams.replace_sampler_ddp is False:
+            return DataLoader(
+                ds,
+                batch_sampler=self.get_custom_sampler(ds),
+                num_workers=self.hparams.dataloader_workers,
+                collate_fn=collate_fn,
+                pin_memory=True,
+            )
+        return DataLoader(
+            ds,
+            batch_size=self.hparams.train_batchsize,
+            num_workers=self.hparams.dataloader_workers,
+            collate_fn=collate_fn,
+            pin_memory=True,
+        )
+
+    def val_dataloader(self):
+        ds = self.datasets[self.hparams.val_datasets_field]
+        collate_fn = self.collate_fn
+        if hasattr(ds, 'collate_fn'):
+            collate_fn = ds.collate_fn
+
+        return DataLoader(
+            ds,
+            batch_size=self.hparams.val_batchsize,
+            shuffle=False,
+            num_workers=self.hparams.dataloader_workers,
+            collate_fn=collate_fn,
+            sampler=DistributedSampler(
+                ds, shuffle=False),
+            pin_memory=True,
+        )
+
+        # return DataLoader(
+        #     ds, shuffle=False, batch_size=self.hparams.val_batchsize, pin_memory=False, collate_fn=collate_fn,
+        # )
+
+    def test_dataloader(self):
+        ds = self.datasets[self.hparams.test_datasets_field]
+
+        collate_fn = self.collate_fn
+        if collate_fn is None and hasattr(ds, 'collater'):
+            collate_fn = ds.collater
+
+        return DataLoader(
+            ds,
+            batch_size=self.hparams.test_batchsize,
+            shuffle=False,
+            num_workers=self.hparams.dataloader_workers,
+            collate_fn=collate_fn,
+            sampler=DistributedSampler(
+                ds, shuffle=False),
+            pin_memory=True,
+        )

fengshen/data/universal_datamodule/universal_sampler.py

@@ -0,0 +1,125 @@
+# coding=utf-8
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Dataloaders."""
+
+
+import torch
+
+
+class PretrainingSampler:
+
+    def __init__(self, total_samples, consumed_samples, micro_batch_size,
+                 data_parallel_rank, data_parallel_size, drop_last=True):
+        # Keep a copy of input params for later use.
+        self.total_samples = total_samples
+        self.consumed_samples = consumed_samples
+        self.micro_batch_size = micro_batch_size
+        self.data_parallel_rank = data_parallel_rank
+        self.micro_batch_times_data_parallel_size = \
+            self.micro_batch_size * data_parallel_size
+        self.drop_last = drop_last
+
+        # Sanity checks.
+        assert self.total_samples > 0, \
+            'no sample to consume: {}'.format(self.total_samples)
+        assert self.consumed_samples < self.total_samples, \
+            'no samples left to consume: {}, {}'.format(self.consumed_samples,
+                                                        self.total_samples)
+        assert self.micro_batch_size > 0
+        assert data_parallel_size > 0
+        assert self.data_parallel_rank < data_parallel_size, \
+            'data_parallel_rank should be smaller than data size: {}, ' \
+            '{}'.format(self.data_parallel_rank, data_parallel_size)
+
+    def __len__(self):
+        return self.total_samples // self.micro_batch_times_data_parallel_size
+
+    def get_start_end_idx(self):
+        start_idx = self.data_parallel_rank * self.micro_batch_size
+        end_idx = start_idx + self.micro_batch_size
+        return start_idx, end_idx
+
+    def __iter__(self):
+        batch = []
+        # Last batch will be dropped if drop_last is not set False
+        for idx in range(self.consumed_samples, self.total_samples):
+            batch.append(idx)
+            if len(batch) == self.micro_batch_times_data_parallel_size:
+                start_idx, end_idx = self.get_start_end_idx()
+                yield batch[start_idx:end_idx]
+                batch = []
+
+        # Check the last partial batch and see drop_last is set
+        if len(batch) > 0 and not self.drop_last:
+            start_idx, end_idx = self.get_start_end_idx()
+            yield batch[start_idx:end_idx]
+
+
+class PretrainingRandomSampler:
+
+    def __init__(self, total_samples, consumed_samples, micro_batch_size,
+                 data_parallel_rank, data_parallel_size, epoch):
+        # Keep a copy of input params for later use.
+        self.total_samples = total_samples
+        self.consumed_samples = consumed_samples
+        self.micro_batch_size = micro_batch_size
+        self.data_parallel_rank = data_parallel_rank
+        self.data_parallel_size = data_parallel_size
+        self.micro_batch_times_data_parallel_size = \
+            self.micro_batch_size * data_parallel_size
+        self.last_batch_size = \
+            self.total_samples % self.micro_batch_times_data_parallel_size
+        self.epoch = epoch
+
+        # Sanity checks.
+        assert self.total_samples > 0, \
+            'no sample to consume: {}'.format(self.total_samples)
+        assert self.micro_batch_size > 0
+        assert data_parallel_size > 0
+        assert self.data_parallel_rank < data_parallel_size, \
+            'data_parallel_rank should be smaller than data size: {}, ' \
+            '{}'.format(self.data_parallel_rank, data_parallel_size)
+
+    def __len__(self):
+        return self.total_samples // self.micro_batch_times_data_parallel_size
+
+    def __iter__(self):
+        active_total_samples = self.total_samples - self.last_batch_size
+        current_epoch_samples = self.consumed_samples % active_total_samples
+        assert current_epoch_samples % self.micro_batch_times_data_parallel_size == 0
+
+        # data sharding and random sampling
+        bucket_size = (self.total_samples // self.micro_batch_times_data_parallel_size) \
+            * self.micro_batch_size
+        bucket_offset = current_epoch_samples // self.data_parallel_size
+        start_idx = self.data_parallel_rank * bucket_size
+
+        g = torch.Generator()
+        g.manual_seed(self.epoch)
+        random_idx = torch.randperm(bucket_size, generator=g).tolist()
+        idx_range = [start_idx + x for x in random_idx[bucket_offset:]]
+
+        batch = []
+        # Last batch if not complete will be dropped.
+        for idx in idx_range:
+            batch.append(idx)
+            if len(batch) == self.micro_batch_size:
+                self.consumed_samples += self.micro_batch_times_data_parallel_size
+                yield batch
+                batch = []
+
+    def set_epoch(self, epoch):
+        self.epoch = epoch

fengshen/examples/DAVAE/generate.py

@@ -0,0 +1,36 @@
+# -*- encoding: utf-8 -*-
+'''
+Copyright 2022 The International Digital Economy Academy (IDEA). CCNL team. All rights reserved.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+@File    :   generate.py
+@Time    :   2022/11/04 19:17
+@Author  :   Liang Yuxin
+@Version :   1.0
+@Contact :   liangyuxin@idea.edu.cn
+@License :   (C)Copyright 2022-2023, CCNL-IDEA
+'''
+# here put the import lib
+
+import torch
+from fengshen.models.DAVAE.DAVAEModel import DAVAEModel
+from transformers import BertTokenizer,T5Tokenizer
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+encoder_tokenizer = BertTokenizer.from_pretrained("IDEA-CCNL/Randeng-DAVAE-1.2B-General-Chinese")
+decoder_tokenizer = T5Tokenizer.from_pretrained("IDEA-CCNL/Randeng-DAVAE-1.2B-General-Chinese", eos_token = '<|endoftext|>', pad_token = '<pad>',extra_ids=0)
+decoder_tokenizer.add_special_tokens({'bos_token':'<bos>'})
+vae_model = DAVAEModel.from_pretrained("IDEA-CCNL/Randeng-DAVAE-1.2B-General-Chinese").to(device)
+input_texts = [
+    "针对电力系统中的混沌振荡对整个互联电网的危害问题,提出了一种基于非线性光滑函数的滑模控制方法.",
+    "超市面积不算大.挺方便附近的居民购买的. 生活用品也比较齐全.价格适用中.",
+]
+output_texts = vae_model.simulate_batch(encoder_tokenizer,decoder_tokenizer,input_texts)
+print(output_texts)

fengshen/examples/FastDemo/README.md

@@ -0,0 +1,105 @@
+# 「streamlit」快速搭建你的算法demo
+在搭建demo之前，首先得做好这些准备工作：
+- 模型训练完毕
+- 模型的入参确定
+- 安装streamlit库，`pip install streamlit` 就可以安装。
+
+streamlit脚本的启动方式是 `streamlit run demo.py`，很简单就启动了一个demo页面，页面会随着脚本代码的改变实时刷新的。所以在没有经验的时候，可以创建一个demo.py的文件，照着下面的教程一步一步添加代码，看页面的展示情况。下面开始上干货，具体细节在代码注释中有说明！
+
+### 第一步 导包
+```python 
+import streamlit as st
+# 其他包更具你的需要导入
+```
+[streamlit](https://streamlit.io)是一个用于构建机器学习、深度学习、数据可视化demo的python框架。它不需要你有web开发的经验，会写python就可以高效的开发你的demo。
+
+### 第二步 页面导航信息以及布局配置
+
+```python 
+st.set_page_config(
+     page_title="余元医疗问答", # 页面标签标题
+     page_icon=":shark:", # 页面标签图标
+     layout="wide", # 页面的布局
+     initial_sidebar_state="expanded", # 左侧的sidebar的布局方式
+     # 配置菜单按钮的信息
+     menu_items={
+         'Get Help': 'https://www.extremelycoolapp.com/help',
+         'Report a bug': "https://www.extremelycoolapp.com/bug",
+         'About': "# This is a header. This is an *extremely* cool app!"
+     }
+ )
+```
+这一步可以省略，如果想让app更加个性化，可以添加这些设置。
+
+### 第三步 设置demo标题
+```python 
+st.title('Demo for MedicalQA') 
+```
+streamlit的每一个小组件对应于页面都有一个默认的样式展示。
+
+### 第四步 配置demo的参数
+
+```python 
+# 此处是用的sidebar，侧边栏作为参数配置模块
+st.sidebar.header("参数配置")
+# 这里是在sidebar里面创建了表单，每个表单一定有一个标题和提交按钮
+sbform = st.sidebar.form("固定参数设置")
+# slider是滑动条组建，可以配置数值型参数
+n_sample = sbform.slider("设置返回条数",min_value=1,max_value=10,value=3)
+text_length = sbform.slider('生成长度:',min_value=32,max_value=512,value=64,step=32)
+text_level = sbform.slider('文本多样性:',min_value=0.1,max_value=1.0,value=0.9,step=0.1)
+# number_input也可以配置数值型参数
+model_id = sbform.number_input('选择模型号:',min_value=0,max_value=13,value=13,step=1)
+# selectbox选择组建，只能选择配置的选项
+trans = sbform.selectbox('选择翻译内核',['百度通用','医疗生物'])
+# 提交表单的配置，这些参数的赋值才生效
+sbform.form_submit_button("提交配置")
+
+# 这里是页面中的参数配置，也是demo的主体之一
+form = st.form("参数设置")
+# 本demo是qa demo，所以要录入用户的文本输入，text_input组建可以实现
+input_text = form.text_input('请输入你的问题:',value='',placeholder='例如：糖尿病的症状有哪些？')
+form.form_submit_button("提交")
+```
+以上就把demo的参数基本配置完成了。
+
+### 第五步 模型预测
+```python 
+# 定义一个前向预测的方法
+# @st.cache(suppress_st_warning=True)
+def generate_qa(input_text,n_sample,model_id='7',length=64,translator='baidu',level=0.7):
+    # 这里我们是把模型用fastapi搭建了一个api服务
+    URL = 'http://192.168.190.63:6605/qa'
+    data = {
+            "text":input_text,"n_sample":n_sample,
+            "model_id":model_id,"length":length,
+            'translator':translator,'level':level
+            }
+    r = requests.get(URL,params=data)
+    return r.text
+# 模型预测结果
+results = generate_qa(input_text,n_sample,model_id=str(model_id),
+                    translator=translator,length=text_length,level=text_level)
+```
+这里说明一下，由于demo展示机器没有GPU，所以模型部署采用的是Fastapi部署在后台的。如果demo展示的机器可以直接部署模型，这里可以直接把模型预测的方法写在这里，不需要另外部署模型，再用api的方式调用。这样做有一个值得注意的地方，因为streamlit的代码每一次运行，都是从头到尾执行一遍，就导致模型可能会重复加载，所以这里需要用到st.cache组建，当内容没有更新的时候，会把这一步的结果缓存，而不会重新执行。保证了效率不会因此而下降。
+
+### 第六步 结果展示
+```python 
+with st.spinner('老夫正在思考中🤔...'):
+    if input_text:
+        results = generate_qa(input_text,n_sample,model_id=str(model_id),
+                        translator=translator,length=text_length,level=text_level)
+        for idx,item in enumerate(eval(results),start=1):
+            st.markdown(f"""
+            **候选回答「{idx}」:**\n
+            """)
+            st.info('中文：%s'%item['fy_next_sentence'])
+            st.info('英文：%s'%item['next_sentence'])
+```
+streamlit对不同格式的内容展示，有丰富的组建，对于文本可以用`st.markdown`组建以及`st.text`和`st.write`展示。更多组建和功能可以参考官方文档：https://docs.streamlit.io
+
+至此，一个完整的demo展示就完成了。效果图如下：
+
+![](./image/demo.png)
+
+完整的代码可以参考：`Fengshenbang-LM/fengshen/examples/FastDemo/YuyuanQA.py`

fengshen/examples/FastDemo/YuyuanQA.py

@@ -0,0 +1,71 @@
+import requests
+import langid
+import streamlit as st
+from translate import baiduTranslatorMedical
+from translate import baiduTranslator
+
+langid.set_languages(['en', 'zh'])
+lang_dic = {'zh': 'en', 'en': 'zh'}
+
+st.set_page_config(
+    page_title="余元医疗问答",
+    page_icon=":shark:",
+    #  layout="wide",
+    initial_sidebar_state="expanded",
+    menu_items={
+        'Get Help': 'https://www.extremelycoolapp.com/help',
+        'Report a bug': "https://www.extremelycoolapp.com/bug",
+        'About': "# This is a header. This is an *extremely* cool app!"
+    }
+)
+st.title('Demo for MedicalQA')
+
+
+st.sidebar.header("参数配置")
+sbform = st.sidebar.form("固定参数设置")
+n_sample = sbform.slider("设置返回条数", min_value=1, max_value=10, value=3)
+text_length = sbform.slider('生成长度:', min_value=32, max_value=512, value=64, step=32)
+text_level = sbform.slider('文本多样性:', min_value=0.1, max_value=1.0, value=0.9, step=0.1)
+model_id = sbform.number_input('选择模型号:', min_value=0, max_value=13, value=13, step=1)
+trans = sbform.selectbox('选择翻译内核', ['百度通用', '医疗生物'])
+sbform.form_submit_button("配置")
+
+
+form = st.form("参数设置")
+input_text = form.text_input('请输入你的问题:', value='', placeholder='例如：糖尿病的症状有哪些？')
+if trans == '百度通用':
+    translator = 'baidu_common'
+else:
+    translator = 'baidu'
+if input_text:
+    lang = langid.classify(input_text)[0]
+    if translator == 'baidu':
+        st.write('**你的问题是:**', baiduTranslatorMedical(input_text, src=lang, dest=lang_dic[lang]).text)
+    else:
+        st.write('**你的问题是:**', baiduTranslator(input_text, src=lang, dest=lang_dic[lang]).text)
+
+form.form_submit_button("提交")
+
+# @st.cache(suppress_st_warning=True)
+
+
+def generate_qa(input_text, n_sample, model_id='7', length=64, translator='baidu', level=0.7):
+    # st.write('调用了generate函数')
+    URL = 'http://192.168.190.63:6605/qa'
+    data = {"text": input_text, "n_sample": n_sample, "model_id": model_id,
+            "length": length, 'translator': translator, 'level': level}
+    r = requests.get(URL, params=data)
+    return r.text
+# my_bar = st.progress(80)
+
+
+with st.spinner('老夫正在思考中🤔...'):
+    if input_text:
+        results = generate_qa(input_text, n_sample, model_id=str(model_id),
+                              translator=translator, length=text_length, level=text_level)
+        for idx, item in enumerate(eval(results), start=1):
+            st.markdown(f"""
+            **候选回答「{idx}」:**\n
+            """)
+            st.info('中文：%s' % item['fy_next_sentence'])
+            st.info('英文：%s' % item['next_sentence'])

fengshen/examples/GAVAE/generate.py

@@ -0,0 +1,23 @@
+import torch
+from transformers import BertTokenizer,T5Tokenizer
+from fengshen.models.GAVAE.GAVAEModel import GAVAEModel
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+encoder_tokenizer = BertTokenizer.from_pretrained("IDEA-CCNL/Randeng-GAVAE-1.2B-Augmentation-Chinese")
+decoder_tokenizer = T5Tokenizer.from_pretrained("IDEA-CCNL/Randeng-GAVAE-1.2B-Augmentation-Chinese", eos_token = '<|endoftext|>', pad_token = '<pad>',extra_ids=0)
+decoder_tokenizer.add_special_tokens({'bos_token':'<bos>'})
+input_texts = [
+    "非常好的一个博物馆，是我所有去过的博物馆里感觉最正规的一家，凭有效证件可以入馆，可以自助免费存小件物品，讲解员和馆内外的工作人员也非常认真，其他的服务人员也很热情，非常好的！馆内的藏品也让人非常震撼！希望继续保持～", 
+    "这是我来长沙最最期待的一定要去的地方，总算今天特地去瞻仰千古遗容了，开车到门口大屏幕显示着门票已发完的字样，心里一惊以为今天是白来了。但进了停车场才知道凭停车卡和有效身份证里面也能领，停车还不花钱，真好。", 
+    "地方很大 很气派~~可以逛很久~~~去的时候是免费的~不过要安检~~~里面的马王堆~幸追夫人~还是很不错的~~~~去的时候有一个吴越文化特别展~~~东西也很多~~~~~很好看",
+    "我们到达的时候是下午3点，门票已经发完了。当时正焦虑的不知道怎么办才好，门卫大哥给我们俩补办了门票，这才得以入馆。非常感谢！绝对不虚此行！相当震撼的展览！原来古人也化妆，还有假发。记忆最深的是那个藕汤。可惜真颜已不得见。", 
+    "去过三次，个人认为这是长沙最值得去的地方，博物馆的重点就是辛追，遗憾的是，每次去我都会感到悲哀，虽然我三次去的时候都要门票，但是每次看到辛追，都觉得现代的人类不应该挖她出来，除了第一次我觉得辛追像刚死去一样，后来两次我觉得太惨不忍睹了。建议大家要去就早去，以后肯定越来越腐烂", 
+    "上大学时候去的，当时学生证是半价25，后来凭有效证件就不要钱了。非常喜欢的一家博物馆，里面可看的东西很多，当然最吸引我的就是那个辛追夫人和“素纱单衣”，果然不是盖的~里面的讲解员大部分都是师大学历史类的，非常专业和有耐心。虽然不在长沙了，不过对那里还是很有感情的，赞~~~", 
+    "这两年也有很多机会去博物馆。。。不过还是想说湖南省博物馆是非常有特色的。。。应该说整个展览分成两个部分吧。。。一个部分是马王堆的主体展。。。另一个就是湖南的一些考古发现。。。其实来省博大部分的游客还是冲着马王堆来的吧。。。博物馆也很有心的为每一批游客安排了讲解员。。。从马王堆的发现到马王堆出土文物的介绍再到最后棺木和辛追的介绍。。。真是上了一节很生动的历史课。",
+    "网上订票去的，还是很顺利的就进去了，里面挺清净的，外围的环境也不错，还有鸽子可以喂。那天不是很闹，兜了一圈感觉还是很顺畅的，老娘娘和金缕玉衣挺震撼的。到此一游还是挺需要的",
+]
+gavae_model = GAVAEModel.from_pretrained("IDEA-CCNL/Randeng-GAVAE-1.2B-Augmentation-Chinese").to(device)
+gavae_model.train_gan(encoder_tokenizer,decoder_tokenizer,input_texts)
+# n:输出样本数量
+texts = gavae_model.generate(n=5)
+print(texts)

fengshen/examples/PPVAE/generate.py

@@ -0,0 +1,24 @@
+import torch
+from transformers import BertTokenizer,T5Tokenizer
+from fengshen.models.PPVAE.pluginVAE import PPVAEModel
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+encoder_tokenizer = BertTokenizer.from_pretrained("IDEA-CCNL/Randeng-PPVAE-1.2B-Augmentation-Chinese")
+decoder_tokenizer = T5Tokenizer.from_pretrained("IDEA-CCNL/Randeng-PPVAE-1.2B-Augmentation-Chinese", eos_token = '<|endoftext|>', pad_token = '<pad>',extra_ids=0)
+decoder_tokenizer.add_special_tokens({'bos_token':'<bos>'})
+ppvae_model = PPVAEModel.from_pretrained("IDEA-CCNL/Randeng-PPVAE-1.2B-Augmentation-Chinese").to(device)
+input_texts = [
+    "非常好的一个博物馆，是我所有去过的博物馆里感觉最正规的一家，凭有效证件可以入馆，可以自助免费存小件物品，讲解员和馆内外的工作人员也非常认真，其他的服务人员也很热情，非常好的！馆内的藏品也让人非常震撼！希望继续保持～", 
+    "这是我来长沙最最期待的一定要去的地方，总算今天特地去瞻仰千古遗容了，开车到门口大屏幕显示着门票已发完的字样，心里一惊以为今天是白来了。但进了停车场才知道凭停车卡和有效身份证里面也能领，停车还不花钱，真好。", 
+    "地方很大 很气派~~可以逛很久~~~去的时候是免费的~不过要安检~~~里面的马王堆~幸追夫人~还是很不错的~~~~去的时候有一个吴越文化特别展~~~东西也很多~~~~~很好看",
+    "我们到达的时候是下午3点，门票已经发完了。当时正焦虑的不知道怎么办才好，门卫大哥给我们俩补办了门票，这才得以入馆。非常感谢！绝对不虚此行！相当震撼的展览！原来古人也化妆，还有假发。记忆最深的是那个藕汤。可惜真颜已不得见。", 
+    "去过三次，个人认为这是长沙最值得去的地方，博物馆的重点就是辛追，遗憾的是，每次去我都会感到悲哀，虽然我三次去的时候都要门票，但是每次看到辛追，都觉得现代的人类不应该挖她出来，除了第一次我觉得辛追像刚死去一样，后来两次我觉得太惨不忍睹了。建议大家要去就早去，以后肯定越来越腐烂", 
+    "上大学时候去的，当时学生证是半价25，后来凭有效证件就不要钱了。非常喜欢的一家博物馆，里面可看的东西很多，当然最吸引我的就是那个辛追夫人和“素纱单衣”，果然不是盖的~里面的讲解员大部分都是师大学历史类的，非常专业和有耐心。虽然不在长沙了，不过对那里还是很有感情的，赞~~~", 
+    "这两年也有很多机会去博物馆。。。不过还是想说湖南省博物馆是非常有特色的。。。应该说整个展览分成两个部分吧。。。一个部分是马王堆的主体展。。。另一个就是湖南的一些考古发现。。。其实来省博大部分的游客还是冲着马王堆来的吧。。。博物馆也很有心的为每一批游客安排了讲解员。。。从马王堆的发现到马王堆出土文物的介绍再到最后棺木和辛追的介绍。。。真是上了一节很生动的历史课。",
+    "网上订票去的，还是很顺利的就进去了，里面挺清净的，外围的环境也不错，还有鸽子可以喂。那天不是很闹，兜了一圈感觉还是很顺畅的，老娘娘和金缕玉衣挺震撼的。到此一游还是挺需要的",
+]
+
+ppvae_model.train_plugin(encoder_tokenizer,decoder_tokenizer,input_texts,negative_samples=None)
+# n:输出样本数量
+texts = ppvae_model.generate(n=5)
+print(texts)

fengshen/examples/classification/demo_classification_afqmc_erlangshen_offload.sh

@@ -0,0 +1,103 @@
+MODEL_NAME="IDEA-CCNL/Erlangshen-MegatronBert-1.3B"
+
+TEXTA_NAME=sentence1
+TEXTB_NAME=sentence2
+LABEL_NAME=label
+ID_NAME=id
+
+BATCH_SIZE=1
+VAL_BATCH_SIZE=1
+ZERO_STAGE=3
+config_json="./ds_config.json"
+
+cat <<EOT > $config_json
+{
+  "train_micro_batch_size_per_gpu": $BATCH_SIZE,
+  "steps_per_print": 1000,
+  "gradient_clipping": 1,
+  "zero_optimization": {
+        "stage": ${ZERO_STAGE},
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "offload_param": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "sub_group_size": 1e9,
+        "stage3_max_live_parameters": 1e9,
+        "stage3_max_reuse_distance": 1e9
+    },
+  "zero_allow_untested_optimizer": false,
+  "fp16": {
+    "enabled": true,
+    "loss_scale": 0,
+    "loss_scale_window": 1000,
+    "hysteresis": 2,
+    "min_loss_scale": 1
+  },
+  "activation_checkpointing": {
+    "partition_activations": false,
+    "contiguous_memory_optimization": false
+  },
+  "wall_clock_breakdown": false
+}
+EOT
+
+export PL_DEEPSPEED_CONFIG_PATH=$config_json
+
+DATA_ARGS="\
+        --dataset_name IDEA-CCNL/AFQMC \
+        --train_batchsize $BATCH_SIZE \
+        --valid_batchsize $VAL_BATCH_SIZE \
+        --max_length 128 \
+        --texta_name $TEXTA_NAME \
+        --textb_name $TEXTB_NAME \
+        --label_name $LABEL_NAME \
+        --id_name $ID_NAME \
+        "
+
+MODEL_ARGS="\
+        --learning_rate 1e-5 \
+        --weight_decay 1e-1 \
+        --warmup_ratio 0.01 \
+        --num_labels 2 \
+        --model_type huggingface-auto \
+        "
+
+MODEL_CHECKPOINT_ARGS="\
+        --monitor val_acc \
+        --save_top_k 3 \
+        --mode max \
+        --every_n_train_steps 0 \
+        --save_weights_only True \
+        --dirpath . \
+        --filename model-{epoch:02d}-{val_acc:.4f} \
+        "
+
+
+TRAINER_ARGS="\
+        --max_epochs 67 \
+        --gpus 1 \
+        --num_nodes 1 \
+        --strategy deepspeed_stage_${ZERO_STAGE}_offload \
+        --gradient_clip_val 1.0 \
+        --check_val_every_n_epoch 1 \
+        --val_check_interval 1.0 \
+        --precision 16 \
+        --default_root_dir . \
+        "
+
+options=" \
+        --pretrained_model_path $MODEL_NAME \
+        $DATA_ARGS \
+        $MODEL_ARGS \
+        $MODEL_CHECKPOINT_ARGS \
+        $TRAINER_ARGS \
+        "
+ 
+python3 finetune_classification.py $options
+