一、出处
https://www.sbert.net/examples/training/sts/README.html
https://github.com/UKPLab/sentence-transformers/blob/master/examples/training/sts/training_stsbenchmark.py
二、代码
此示例从头开始为 STSbenchmark 训练 BERT(或任何其他转换器模型,如 RoBERTa、DistilBERT 等)。 它生成句子嵌入,可以使用余弦相似度进行比较以测量相似度。
用法:
python training_nli.py 或者 python training_nli.py pretrained_transformer_model_name
from torch.utils.data import DataLoader import math from sentence_transformers import SentenceTransformer, LoggingHandler, losses, models, util from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator from sentence_transformers.readers import InputExample import logging from datetime import datetime import sys import os import gzip import csv #### Just some code to print debug information to stdout
只是一些将调试信息打印到标准输出的代码
logging.basicConfig(format=‘%(asctime)s - %(message)s‘, datefmt=‘%Y-%m-%d %H:%M:%S‘, level=logging.INFO, handlers=[LoggingHandler()]) #### /print debug information to stdout 调试信息打印到标准输出 #Check if dataset exsist. If not, download and extract it 检查数据集是否存在。 如果没有,请下载并解压 sts_dataset_path = ‘datasets/stsbenchmark.tsv.gz‘ if not os.path.exists(sts_dataset_path): util.http_get(‘https://sbert.net/datasets/stsbenchmark.tsv.gz‘, sts_dataset_path) #You can specify any huggingface/transformers pre-trained model here, for example, bert-base-uncased, roberta-base, xlm-roberta-base
您可以在此处指定任何 Huggingface/transformers 预训练模型,例如,bert-base-uncased、roberta-base、xlm-roberta-base
model_name = sys.argv[1] if len(sys.argv) > 1 else ‘distilbert-base-uncased‘ # Read the dataset 读取数据集 train_batch_size = 16 num_epochs = 4 model_save_path = ‘output/training_stsbenchmark_‘+model_name.replace("/", "-")+‘-‘+datetime.now().strftime("%Y-%m-%d_%H-%M-%S") # Use Huggingface/transformers model (like BERT, RoBERTa, XLNet, XLM-R) for mapping tokens to embeddings
使用 Huggingface/transformers 模型(如 BERT、RoBERTa、XLNet、XLM-R)将令牌映射到嵌入
word_embedding_model = models.Transformer(model_name) # Apply mean pooling to get one fixed sized sentence vector 应用平均池化得到一个固定大小的句子向量 pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension(), pooling_mode_mean_tokens=True, pooling_mode_cls_token=False, pooling_mode_max_tokens=False) model = SentenceTransformer(modules=[word_embedding_model, pooling_model]) # Convert the dataset to a DataLoader ready for training 将数据集转换为准备训练的 DataLoader logging.info("Read STSbenchmark train dataset") train_samples = [] dev_samples = [] test_samples = [] with gzip.open(sts_dataset_path, ‘rt‘, encoding=‘utf8‘) as fIn: reader = csv.DictReader(fIn, delimiter=‘\t‘, quoting=csv.QUOTE_NONE) for row in reader: score = float(row[‘score‘]) / 5.0 # Normalize score to range 0 ... 1 inp_example = InputExample(texts=[row[‘sentence1‘], row[‘sentence2‘]], label=score) if row[‘split‘] == ‘dev‘: dev_samples.append(inp_example) elif row[‘split‘] == ‘test‘: test_samples.append(inp_example) else: train_samples.append(inp_example) train_dataloader = DataLoader(train_samples, shuffle=True, batch_size=train_batch_size) train_loss = losses.CosineSimilarityLoss(model=model) logging.info("Read STSbenchmark dev dataset") evaluator = EmbeddingSimilarityEvaluator.from_input_examples(dev_samples, name=‘sts-dev‘) # Configure the training. We skip evaluation in this example 配置训练。 我们在这个例子中跳过评估 warmup_steps = math.ceil(len(train_dataloader) * num_epochs * 0.1) #10% of train data for warm-up logging.info("Warmup-steps: {}".format(warmup_steps)) # Train the model model.fit(train_objectives=[(train_dataloader, train_loss)], evaluator=evaluator, epochs=num_epochs, evaluation_steps=1000, warmup_steps=warmup_steps, output_path=model_save_path) ############################################################################## # # Load the stored model and evaluate its performance on STS benchmark dataset 加载存储的模型并评估其在 STS 基准数据集上的性能 # ############################################################################## model = SentenceTransformer(model_save_path) test_evaluator = EmbeddingSimilarityEvaluator.from_input_examples(test_samples, name=‘sts-test‘) test_evaluator(model, output_path=model_save_path)
三、评估
""" This examples loads a pre-trained model and evaluates it on the STSbenchmark dataset
此示例加载预训练模型并在 STSbenchmark 数据集上对其进行评估 Usage: python evaluation_stsbenchmark.py OR python evaluation_stsbenchmark.py model_name """ from sentence_transformers import SentenceTransformer, util, LoggingHandler, InputExample from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator import logging import sys import torch import gzip import os import csv script_folder_path = os.path.dirname(os.path.realpath(__file__)) #Limit torch to 4 threads 将割炬限制为 4 个线程 torch.set_num_threads(4) #### Just some code to print debug information to stdout logging.basicConfig(format=‘%(asctime)s - %(message)s‘, datefmt=‘%Y-%m-%d %H:%M:%S‘, level=logging.INFO, handlers=[LoggingHandler()]) #### /print debug information to stdout model_name = sys.argv[1] if len(sys.argv) > 1 else ‘stsb-distilroberta-base-v2‘ # Load a named sentence model (based on BERT). This will download the model from our server. # Alternatively, you can also pass a filepath to SentenceTransformer()
加载命名句子模型(基于 BERT)。 这将从我们的服务器下载模型。 或者,您也可以将文件路径传递给 SentenceTransformer()
model = SentenceTransformer(model_name) sts_dataset_path = ‘data/stsbenchmark.tsv.gz‘ if not os.path.exists(sts_dataset_path): util.http_get(‘https://sbert.net/datasets/stsbenchmark.tsv.gz‘, sts_dataset_path) train_samples = [] dev_samples = [] test_samples = [] with gzip.open(sts_dataset_path, ‘rt‘, encoding=‘utf8‘) as fIn: reader = csv.DictReader(fIn, delimiter=‘\t‘, quoting=csv.QUOTE_NONE) for row in reader: score = float(row[‘score‘]) / 5.0 # Normalize score to range 0 ... 1 inp_example = InputExample(texts=[row[‘sentence1‘], row[‘sentence2‘]], label=score) if row[‘split‘] == ‘dev‘: dev_samples.append(inp_example) elif row[‘split‘] == ‘test‘: test_samples.append(inp_example) else: train_samples.append(inp_example) evaluator = EmbeddingSimilarityEvaluator.from_input_examples(dev_samples, name=‘sts-dev‘) model.evaluate(evaluator) evaluator = EmbeddingSimilarityEvaluator.from_input_examples(test_samples, name=‘sts-test‘) model.evaluate(evaluator)