""" This script provides an example to wrap TencentPretrain for SimCSE. """ import sys import os import random import argparse import math import scipy.stats import torch import torch.nn as nn import numpy as np tencentpretrain_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) sys.path.append(tencentpretrain_dir) from tencentpretrain.embeddings import * from tencentpretrain.encoders import * from tencentpretrain.targets import * from tencentpretrain.utils.vocab import Vocab from tencentpretrain.utils.constants import * from tencentpretrain.utils import * from tencentpretrain.utils.optimizers import * from tencentpretrain.utils.config import load_hyperparam from tencentpretrain.utils.seed import set_seed from tencentpretrain.utils.logging import init_logger from tencentpretrain.model_saver import save_model from tencentpretrain.opts import finetune_opts, tokenizer_opts from finetune.run_classifier import count_labels_num, build_optimizer, load_or_initialize_parameters from finetune.run_classifier_siamese import batch_loader class SimCSE(nn.Module): def __init__(self, args): super(SimCSE, self).__init__() self.embedding = Embedding(args) for embedding_name in args.embedding: tmp_emb = str2embedding[embedding_name](args, len(args.tokenizer.vocab)) self.embedding.update(tmp_emb, embedding_name) self.encoder = str2encoder[args.encoder](args) self.pooling_type = args.pooling def forward(self, src, seg): """ Args: src: [batch_size x seq_length] tgt: [batch_size] seg: [batch_size x seq_length] """ # Embedding. emb_0 = self.embedding(src[0], seg[0]) emb_1 = self.embedding(src[1], seg[1]) # Encoder. output_0 = self.encoder(emb_0, seg[0]) output_1 = self.encoder(emb_1, seg[1]) # Target. features_0 = self.pooling(output_0, seg[0], self.pooling_type) features_1 = self.pooling(output_1, seg[1], self.pooling_type) return features_0, features_1 def pooling(self, memory_bank, seg, pooling_type): seg = torch.unsqueeze(seg, dim=-1).type(torch.float) memory_bank = memory_bank * seg if pooling_type == "mean": features = torch.sum(memory_bank, dim=1) features = torch.div(features, torch.sum(seg, dim=1)) elif pooling_type == "last": features = memory_bank[torch.arange(memory_bank.shape[0]), torch.squeeze(torch.sum(seg, dim=1).type(torch.int64) - 1), :] elif pooling_type == "max": features = torch.max(memory_bank + (seg - 1) * sys.maxsize, dim=1)[0] else: features = memory_bank[:, 0, :] return features def read_dataset(args, path): dataset, columns = [], {} with open(path, mode="r", encoding="utf-8") as f: for line_id, line in enumerate(f): if line_id == 0: for i, column_name in enumerate(line.rstrip("\r\n").split("\t")): columns[column_name] = i continue line = line.rstrip("\r\n").split("\t") if "text_b" in columns: text_a, text_b = line[columns["text_a"]], line[columns["text_b"]] else: text_a = line[columns["text_a"]] text_b = text_a src_a = args.tokenizer.convert_tokens_to_ids([CLS_TOKEN] + args.tokenizer.tokenize(text_a) + [SEP_TOKEN]) src_b = args.tokenizer.convert_tokens_to_ids([CLS_TOKEN] + args.tokenizer.tokenize(text_b) + [SEP_TOKEN]) seg_a = [1] * len(src_a) seg_b = [1] * len(src_b) PAD_ID = args.tokenizer.convert_tokens_to_ids([PAD_TOKEN])[0] if len(src_a) >= args.seq_length: src_a = src_a[:args.seq_length] seg_a = seg_a[:args.seq_length] while len(src_a) < args.seq_length: src_a.append(PAD_ID) seg_a.append(0) if len(src_b) >= args.seq_length: src_b = src_b[:args.seq_length] seg_b = seg_b[:args.seq_length] while len(src_b) < args.seq_length: src_b.append(PAD_ID) seg_b.append(0) if "label" in columns: tgt = float(line[columns["label"]]) dataset.append(((src_a, src_b), tgt, (seg_a, seg_b))) else: dataset.append(((src_a, src_a), -1, (seg_a, seg_a))) return dataset def evaluate(args, dataset): src_a = torch.LongTensor([example[0][0] for example in dataset]) src_b = torch.LongTensor([example[0][1] for example in dataset]) tgt = torch.FloatTensor([example[1] for example in dataset]) seg_a = torch.LongTensor([example[2][0] for example in dataset]) seg_b = torch.LongTensor([example[2][1] for example in dataset]) all_similarities = [] batch_size = args.batch_size args.model.eval() for i, (src_batch, tgt_batch, seg_batch) in enumerate(batch_loader(batch_size, (src_a, src_b), tgt, (seg_a, seg_b))): src_a_batch, src_b_batch = src_batch seg_a_batch, seg_b_batch = seg_batch src_a_batch = src_a_batch.to(args.device) src_b_batch = src_b_batch.to(args.device) seg_a_batch = seg_a_batch.to(args.device) seg_b_batch = seg_b_batch.to(args.device) with torch.no_grad(): features_0, features_1 = args.model((src_a_batch, src_b_batch), (seg_a_batch, seg_b_batch)) similarity_matrix = similarity(features_0, features_1, 1) for j in range(similarity_matrix.size(0)): all_similarities.append(similarity_matrix[j][j].item()) corrcoef = scipy.stats.spearmanr(tgt, all_similarities).correlation args.logger.info("Spearman's correlation: {:.4f}".format(corrcoef)) return corrcoef def similarity(x, y, temperature): x = x / x.norm(dim=-1, keepdim=True) y = y / y.norm(dim=-1, keepdim=True) return torch.matmul(x, y.transpose(-2, -1)) / temperature def main(): parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) finetune_opts(parser) tokenizer_opts(parser) parser.add_argument("--temperature", type=float, default=0.05) parser.add_argument("--eval_steps", type=int, default=200, help="Evaluate frequency.") args = parser.parse_args() # Load the hyperparameters from the config file. args = load_hyperparam(args) set_seed(args.seed) # Build tokenizer. args.tokenizer = str2tokenizer[args.tokenizer](args) # Build classification model. model = SimCSE(args) # Load or initialize parameters. load_or_initialize_parameters(args, model) # Get logger. args.logger = init_logger(args) args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = model.to(args.device) # Training phase. trainset = read_dataset(args, args.train_path) instances_num = len(trainset) batch_size = args.batch_size args.train_steps = int(instances_num * args.epochs_num / batch_size) + 1 args.logger.info("Batch size: {}".format(batch_size)) args.logger.info("The number of training instances: {}".format(instances_num)) optimizer, scheduler = build_optimizer(args, model) if args.fp16: try: from apex import amp except ImportError: raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.") model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level) args.amp = amp if torch.cuda.device_count() > 1: args.logger.info("{} GPUs are available. Let's use them.".format(torch.cuda.device_count())) model = torch.nn.DataParallel(model) args.model = model total_loss, result, best_result = 0.0, 0.0, 0.0 args.logger.info("Start training.") for epoch in range(1, args.epochs_num + 1): random.shuffle(trainset) src_a = torch.LongTensor([example[0][0] for example in trainset]) src_b = torch.LongTensor([example[0][1] for example in trainset]) tgt = torch.FloatTensor([example[1] for example in trainset]) seg_a = torch.LongTensor([example[2][0] for example in trainset]) seg_b = torch.LongTensor([example[2][1] for example in trainset]) model.train() for i, (src_batch, tgt_batch, seg_batch) in enumerate(batch_loader(batch_size, (src_a, src_b), tgt, (seg_a, seg_b))): model.zero_grad() src_a_batch, src_b_batch = src_batch seg_a_batch, seg_b_batch = seg_batch src_a_batch = src_a_batch.to(args.device) src_b_batch = src_b_batch.to(args.device) seg_a_batch = seg_a_batch.to(args.device) seg_b_batch = seg_b_batch.to(args.device) features_0, features_1 = model((src_a_batch, src_b_batch), (seg_a_batch, seg_b_batch)) similarity_matrix = similarity(features_0, features_1, args.temperature) tgt_batch = torch.arange(similarity_matrix.size(0), device=similarity_matrix.device, dtype=torch.long) loss = nn.CrossEntropyLoss()(similarity_matrix, tgt_batch) if args.fp16: with args.amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward() else: loss.backward() optimizer.step() scheduler.step() total_loss += loss.item() if (i + 1) % args.report_steps == 0: args.logger.info("Epoch id: {}, Training steps: {}, Avg loss: {:.3f}" .format(epoch, i + 1, total_loss / args.report_steps)) total_loss = 0.0 if (i + 1) % args.eval_steps == 0 or (i + 1) == math.ceil(instances_num / batch_size): result = evaluate(args, read_dataset(args, args.dev_path)) args.logger.info("Epoch id: {}, Training steps: {}, Evaluate result: {}, Best result: {}" .format(epoch, i + 1, result, best_result)) if result > best_result: best_result = result save_model(model, args.output_model_path) args.logger.info("It is the best model until now. Save it to {}".format(args.output_model_path)) if __name__ == "__main__": main()