def __init__(self,
model_path: str,
model_file: str,
vocab_file: str = "vocab.txt",
config_file: str = "config.json",
device: typing.Union[torch.device, int, str] = "cpu",
turns: int = 3,
uttr_len: int = 20,
resp_len: int = 20,
data_type: str = "uru",
add_special_tokens: bool = True):
"""
:param model_path: str型，模型的路径，需要包含训练好的模型，预训练模型pytorch_model.bin以及对应的词表和config文件
:param model_file: str型，训练好的模型的文件名
:param vocab_file: str型，词表的文件名
:param config_file: str型，预训练模型的配置文件
"""
super().__init__()
self.preprocessor = CNAlbertPreprocessorForMultiQA(
Path(model_path) / vocab_file,
uttr_len=uttr_len,
resp_len=resp_len,
add_special_tokens=add_special_tokens)
# 初始化模型
config = AlbertConfig.from_pretrained(Path(model_path) / config_file)
model = AlbertIMN(uttr_len,
resp_len,
turns,
config,
model_path,
data_type=data_type)
cls_task = tasks.Classification(num_classes=2,
losses=[nn.CrossEntropyLoss()])
cls_task.metrics = ['accuracy']
params = model.get_default_params()
params['task'] = cls_task
model.params = params
model.build()
model = model.float()
# 加载模型
self.predictor = Predictor(model,
save_dir=model_path,
checkpoint=model_file,
device=device)
self.device = self.predictor._device
self.turns = turns
self.uttr_len = uttr_len
self.resp_len = resp_len
self.data_type = data_type
from snlp.datagen.dataset.pair_dataset import PairDataset
from snlp.callbacks.padding import MultiQAPadding
from snlp.datagen.dataloader.dict_dataloader import DictDataLoader
from snlp.models.retrieval.dam import DAM
from snlp.models.retrieval.imn import IMN
from snlp.models.retrieval.msn import MSN
from snlp import tasks, metrics, losses
import snlp.tools.constants as constants
from snlp.preprocessors.chinese_preprocessor import CNPreprocessorForMultiQA
from snlp.trainers.predictor import Predictor
os.environ['CUDA_VISIBLE_DEVICES'] = "5, 6, 7"
# 定义任务类型
start = time.time()
cls_task = tasks.Classification(num_classes=2, losses=nn.CrossEntropyLoss())
cls_task.metrics = ['accuracy']
fixed_length_uttr = 20
fixed_length_resp = 20
fixed_length_turn = 5
# 定义所有的模型
MODELS = {
'dam':
DAM(uttr_len=fixed_length_uttr,
resp_len=fixed_length_resp,
turns=fixed_length_turn),
'msn':
MSN(uttr_len=fixed_length_uttr,
resp_len=fixed_length_resp,
from snlp import tasks, metrics, losses
import snlp.tools.constants as constants
from snlp.tools.log import logger
from snlp.preprocessors.chinese_preprocessor import CNAlbertPreprocessorForMultiQA
from snlp.tools.common import seed_everything
from snlp.optimizer import RAdam, AdamW
from snlp.schedule import get_linear_schedule_with_warmup
from snlp.trainers import Trainer
from snlp.losses import LabelSmoothLoss
os.environ['CUDA_VISIBLE_DEVICES'] = "3,4"
# 定义任务类型
start = time.time()
cls_task = tasks.Classification(
num_classes=2,
losses=[nn.CrossEntropyLoss(),
LabelSmoothLoss(label_smoothing=0.2)])
cls_task.metrics = ['accuracy']
fixed_length_uttr = 40
# fixed_length_uttr = 20
fixed_length_resp = 20
fixed_length_turn = 3
name = 'imn'
batch_size = 64
seed = 2020
# lr = 5e-4
# bert_lr = 2e-5
lr = 5e-5
bert_lr = 1e-5