NLP(三十三):sentence-transformers句子相似度官方示例

一、出处

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)

 

NLP(三十三):sentence-transformers句子相似度官方示例

上一篇:07-条件查询


下一篇:DP方案数