初始化项目,由ModelHub XC社区提供模型
Model: vocab-transformers/dense_encoder-msmarco-distilbert-word2vec256k Source: Original Platform
This commit is contained in:
27
.gitattributes
vendored
Normal file
27
.gitattributes
vendored
Normal file
@@ -0,0 +1,27 @@
|
|||||||
|
*.7z filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.arrow filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.bin filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.bin.* filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.ftz filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.gz filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.h5 filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.joblib filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.model filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.onnx filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.ot filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.parquet filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.pb filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.pt filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.pth filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.rar filter=lfs diff=lfs merge=lfs -text
|
||||||
|
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.tar.* filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.tflite filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.tgz filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.xz filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.zip filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.zstandard filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
||||||
7
1_Pooling/config.json
Executable file
7
1_Pooling/config.json
Executable file
@@ -0,0 +1,7 @@
|
|||||||
|
{
|
||||||
|
"word_embedding_dimension": 768,
|
||||||
|
"pooling_mode_cls_token": false,
|
||||||
|
"pooling_mode_mean_tokens": true,
|
||||||
|
"pooling_mode_max_tokens": false,
|
||||||
|
"pooling_mode_mean_sqrt_len_tokens": false
|
||||||
|
}
|
||||||
133
README.md
Executable file
133
README.md
Executable file
@@ -0,0 +1,133 @@
|
|||||||
|
---
|
||||||
|
pipeline_tag: sentence-similarity
|
||||||
|
tags:
|
||||||
|
- sentence-transformers
|
||||||
|
- feature-extraction
|
||||||
|
- sentence-similarity
|
||||||
|
- transformers
|
||||||
|
---
|
||||||
|
|
||||||
|
# dense_encoder-msmarco-distilbert-word2vec256k
|
||||||
|
|
||||||
|
This model is based on [msmarco-word2vec256000-distilbert-base-uncased](https://huggingface.co/nicoladecao/msmarco-word2vec256000-distilbert-base-uncased) with a 256k sized vocabulary initialized with word2vec.
|
||||||
|
|
||||||
|
It has been trained on MS MARCO using [MarginMSELoss](https://github.com/UKPLab/sentence-transformers/blob/master/examples/training/ms_marco/train_bi-encoder_margin-mse.py). See the train_script.py in this repository.
|
||||||
|
Performance:
|
||||||
|
- MS MARCO dev: - (MRR@10)
|
||||||
|
- TREC-DL 2019: 65.53 (nDCG@10)
|
||||||
|
- TREC-DL 2020: 67.42 (nDCG@10)
|
||||||
|
- Avg. on 4 BEIR datasets: 38.97
|
||||||
|
|
||||||
|
|
||||||
|
The word embedding matrix has been frozen while training.
|
||||||
|
|
||||||
|
## Usage (Sentence-Transformers)
|
||||||
|
|
||||||
|
Using this model becomes easy when you have [sentence-transformers](https://www.SBERT.net) installed:
|
||||||
|
|
||||||
|
```
|
||||||
|
pip install -U sentence-transformers
|
||||||
|
```
|
||||||
|
|
||||||
|
Then you can use the model like this:
|
||||||
|
|
||||||
|
```python
|
||||||
|
from sentence_transformers import SentenceTransformer
|
||||||
|
sentences = ["This is an example sentence", "Each sentence is converted"]
|
||||||
|
|
||||||
|
model = SentenceTransformer('{MODEL_NAME}')
|
||||||
|
embeddings = model.encode(sentences)
|
||||||
|
print(embeddings)
|
||||||
|
```
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
## Usage (HuggingFace Transformers)
|
||||||
|
Without [sentence-transformers](https://www.SBERT.net), you can use the model like this: First, you pass your input through the transformer model, then you have to apply the right pooling-operation on-top of the contextualized word embeddings.
|
||||||
|
|
||||||
|
```python
|
||||||
|
from transformers import AutoTokenizer, AutoModel
|
||||||
|
import torch
|
||||||
|
|
||||||
|
|
||||||
|
#Mean Pooling - Take attention mask into account for correct averaging
|
||||||
|
def mean_pooling(model_output, attention_mask):
|
||||||
|
token_embeddings = model_output[0] #First element of model_output contains all token embeddings
|
||||||
|
input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
|
||||||
|
return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
|
||||||
|
|
||||||
|
|
||||||
|
# Sentences we want sentence embeddings for
|
||||||
|
sentences = ['This is an example sentence', 'Each sentence is converted']
|
||||||
|
|
||||||
|
# Load model from HuggingFace Hub
|
||||||
|
tokenizer = AutoTokenizer.from_pretrained('{MODEL_NAME}')
|
||||||
|
model = AutoModel.from_pretrained('{MODEL_NAME}')
|
||||||
|
|
||||||
|
# Tokenize sentences
|
||||||
|
encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
|
||||||
|
|
||||||
|
# Compute token embeddings
|
||||||
|
with torch.no_grad():
|
||||||
|
model_output = model(**encoded_input)
|
||||||
|
|
||||||
|
# Perform pooling. In this case, mean pooling.
|
||||||
|
sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])
|
||||||
|
|
||||||
|
print("Sentence embeddings:")
|
||||||
|
print(sentence_embeddings)
|
||||||
|
```
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
## Evaluation Results
|
||||||
|
|
||||||
|
<!--- Describe how your model was evaluated -->
|
||||||
|
|
||||||
|
For an automated evaluation of this model, see the *Sentence Embeddings Benchmark*: [https://seb.sbert.net](https://seb.sbert.net?model_name={MODEL_NAME})
|
||||||
|
|
||||||
|
|
||||||
|
## Training
|
||||||
|
The model was trained with the parameters:
|
||||||
|
|
||||||
|
**DataLoader**:
|
||||||
|
|
||||||
|
`torch.utils.data.dataloader.DataLoader` of length 7858 with parameters:
|
||||||
|
```
|
||||||
|
{'batch_size': 64, 'sampler': 'torch.utils.data.sampler.RandomSampler', 'batch_sampler': 'torch.utils.data.sampler.BatchSampler'}
|
||||||
|
```
|
||||||
|
|
||||||
|
**Loss**:
|
||||||
|
|
||||||
|
`sentence_transformers.losses.MarginMSELoss.MarginMSELoss`
|
||||||
|
|
||||||
|
Parameters of the fit()-Method:
|
||||||
|
```
|
||||||
|
{
|
||||||
|
"epochs": 30,
|
||||||
|
"evaluation_steps": 0,
|
||||||
|
"evaluator": "NoneType",
|
||||||
|
"max_grad_norm": 1,
|
||||||
|
"optimizer_class": "<class 'transformers.optimization.AdamW'>",
|
||||||
|
"optimizer_params": {
|
||||||
|
"lr": 2e-05
|
||||||
|
},
|
||||||
|
"scheduler": "WarmupLinear",
|
||||||
|
"steps_per_epoch": null,
|
||||||
|
"warmup_steps": 1000,
|
||||||
|
"weight_decay": 0.01
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
|
||||||
|
## Full Model Architecture
|
||||||
|
```
|
||||||
|
SentenceTransformer(
|
||||||
|
(0): Transformer({'max_seq_length': 250, 'do_lower_case': False}) with Transformer model: DistilBertModel
|
||||||
|
(1): Pooling({'word_embedding_dimension': 768, 'pooling_mode_cls_token': False, 'pooling_mode_mean_tokens': True, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False})
|
||||||
|
)
|
||||||
|
```
|
||||||
|
|
||||||
|
## Citing & Authors
|
||||||
|
|
||||||
|
<!--- Describe where people can find more information -->
|
||||||
24
config.json
Executable file
24
config.json
Executable file
@@ -0,0 +1,24 @@
|
|||||||
|
{
|
||||||
|
"_name_or_path": "nicoladecao/msmarco-word2vec256000-distilbert-base-uncased",
|
||||||
|
"activation": "gelu",
|
||||||
|
"architectures": [
|
||||||
|
"DistilBertModel"
|
||||||
|
],
|
||||||
|
"attention_dropout": 0.1,
|
||||||
|
"dim": 768,
|
||||||
|
"dropout": 0.1,
|
||||||
|
"hidden_dim": 3072,
|
||||||
|
"initializer_range": 0.02,
|
||||||
|
"max_position_embeddings": 512,
|
||||||
|
"model_type": "distilbert",
|
||||||
|
"n_heads": 12,
|
||||||
|
"n_layers": 6,
|
||||||
|
"pad_token_id": 0,
|
||||||
|
"qa_dropout": 0.1,
|
||||||
|
"seq_classif_dropout": 0.2,
|
||||||
|
"sinusoidal_pos_embds": false,
|
||||||
|
"tie_weights_": true,
|
||||||
|
"torch_dtype": "float32",
|
||||||
|
"transformers_version": "4.16.2",
|
||||||
|
"vocab_size": 256000
|
||||||
|
}
|
||||||
7
config_sentence_transformers.json
Executable file
7
config_sentence_transformers.json
Executable file
@@ -0,0 +1,7 @@
|
|||||||
|
{
|
||||||
|
"__version__": {
|
||||||
|
"sentence_transformers": "2.2.0",
|
||||||
|
"transformers": "4.16.2",
|
||||||
|
"pytorch": "1.10.2"
|
||||||
|
}
|
||||||
|
}
|
||||||
14
modules.json
Executable file
14
modules.json
Executable file
@@ -0,0 +1,14 @@
|
|||||||
|
[
|
||||||
|
{
|
||||||
|
"idx": 0,
|
||||||
|
"name": "0",
|
||||||
|
"path": "",
|
||||||
|
"type": "sentence_transformers.models.Transformer"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"idx": 1,
|
||||||
|
"name": "1",
|
||||||
|
"path": "1_Pooling",
|
||||||
|
"type": "sentence_transformers.models.Pooling"
|
||||||
|
}
|
||||||
|
]
|
||||||
3
pytorch_model.bin
Executable file
3
pytorch_model.bin
Executable file
@@ -0,0 +1,3 @@
|
|||||||
|
version https://git-lfs.github.com/spec/v1
|
||||||
|
oid sha256:1ab74f1d9ca75b60a82a718c7f60e9dd840656ad718e8748c03ae8f8c8d8e80c
|
||||||
|
size 958156601
|
||||||
4
sentence_bert_config.json
Executable file
4
sentence_bert_config.json
Executable file
@@ -0,0 +1,4 @@
|
|||||||
|
{
|
||||||
|
"max_seq_length": 250,
|
||||||
|
"do_lower_case": false
|
||||||
|
}
|
||||||
1
special_tokens_map.json
Executable file
1
special_tokens_map.json
Executable file
@@ -0,0 +1 @@
|
|||||||
|
{"unk_token": "[UNK]", "sep_token": "[SEP]", "pad_token": "[PAD]", "cls_token": "[CLS]", "mask_token": "[MASK]"}
|
||||||
256165
tokenizer.json
Executable file
256165
tokenizer.json
Executable file
File diff suppressed because it is too large
Load Diff
1
tokenizer_config.json
Executable file
1
tokenizer_config.json
Executable file
@@ -0,0 +1 @@
|
|||||||
|
{"model_max_length": 512, "unk_token": "[UNK]", "cls_token": "[CLS]", "sep_token": "[SEP]", "pad_token": "[PAD]", "mask_token": "[MASK]", "model_input_names": ["input_ids", "attention_mask"], "special_tokens_map_file": "/root/.cache/huggingface/transformers/fe09c361189d8238b9e387f10a088e93f70620bfe74b82036baff1fed512a153.dd8bd9bfd3664b530ea4e645105f557769387b3da9f79bdb55ed556bdd80611d", "name_or_path": "nicoladecao/msmarco-word2vec256000-distilbert-base-uncased", "tokenizer_class": "PreTrainedTokenizerFast"}
|
||||||
235
train_script.py
Executable file
235
train_script.py
Executable file
@@ -0,0 +1,235 @@
|
|||||||
|
import sys
|
||||||
|
import json
|
||||||
|
from torch.utils.data import DataLoader
|
||||||
|
from sentence_transformers import SentenceTransformer, LoggingHandler, util, models, evaluation, losses, InputExample
|
||||||
|
import logging
|
||||||
|
from datetime import datetime
|
||||||
|
import gzip
|
||||||
|
import os
|
||||||
|
import tarfile
|
||||||
|
import tqdm
|
||||||
|
from torch.utils.data import Dataset
|
||||||
|
import random
|
||||||
|
from shutil import copyfile
|
||||||
|
import pickle
|
||||||
|
import argparse
|
||||||
|
|
||||||
|
#### 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
|
||||||
|
|
||||||
|
|
||||||
|
parser = argparse.ArgumentParser()
|
||||||
|
parser.add_argument("--train_batch_size", default=64, type=int)
|
||||||
|
parser.add_argument("--max_seq_length", default=250, type=int)
|
||||||
|
parser.add_argument("--model_name", default="nicoladecao/msmarco-word2vec256000-distilbert-base-uncased")
|
||||||
|
parser.add_argument("--max_passages", default=0, type=int)
|
||||||
|
parser.add_argument("--epochs", default=30, type=int)
|
||||||
|
parser.add_argument("--pooling", default="mean")
|
||||||
|
parser.add_argument("--negs_to_use", default=None, help="From which systems should negatives be used? Multiple systems seperated by comma. None = all")
|
||||||
|
parser.add_argument("--warmup_steps", default=1000, type=int)
|
||||||
|
parser.add_argument("--lr", default=2e-5, type=float)
|
||||||
|
parser.add_argument("--num_negs_per_system", default=5, type=int)
|
||||||
|
parser.add_argument("--use_all_queries", default=False, action="store_true")
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
logging.info(str(args))
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
# The model we want to fine-tune
|
||||||
|
train_batch_size = args.train_batch_size #Increasing the train batch size improves the model performance, but requires more GPU memory
|
||||||
|
model_name = args.model_name
|
||||||
|
max_passages = args.max_passages
|
||||||
|
max_seq_length = args.max_seq_length #Max length for passages. Increasing it, requires more GPU memory
|
||||||
|
|
||||||
|
num_negs_per_system = args.num_negs_per_system # We used different systems to mine hard negatives. Number of hard negatives to add from each system
|
||||||
|
num_epochs = args.epochs # Number of epochs we want to train
|
||||||
|
|
||||||
|
# Load our embedding model
|
||||||
|
|
||||||
|
logging.info("Create new SBERT model")
|
||||||
|
word_embedding_model = models.Transformer(model_name, max_seq_length=max_seq_length)
|
||||||
|
pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension(), args.pooling)
|
||||||
|
model = SentenceTransformer(modules=[word_embedding_model, pooling_model])
|
||||||
|
|
||||||
|
#Freeze embedding layer
|
||||||
|
#word_embedding_model.auto_model.embeddings.requires_grad = False
|
||||||
|
word_embedding_model.auto_model.embeddings.requires_grad_(False)
|
||||||
|
|
||||||
|
model_save_path = f'output-dense/{model_name.replace("/", "-")}-batch_size_{train_batch_size}-{datetime.now().strftime("%Y-%m-%d_%H-%M-%S")}'
|
||||||
|
|
||||||
|
|
||||||
|
# Write self to path
|
||||||
|
os.makedirs(model_save_path, exist_ok=True)
|
||||||
|
|
||||||
|
train_script_path = os.path.join(model_save_path, 'train_script.py')
|
||||||
|
copyfile(__file__, train_script_path)
|
||||||
|
with open(train_script_path, 'a') as fOut:
|
||||||
|
fOut.write("\n\n# Script was called via:\n#python " + " ".join(sys.argv))
|
||||||
|
|
||||||
|
|
||||||
|
### Now we read the MS Marco dataset
|
||||||
|
data_folder = 'msmarco-data'
|
||||||
|
|
||||||
|
#### Read the corpus files, that contain all the passages. Store them in the corpus dict
|
||||||
|
corpus = {} #dict in the format: passage_id -> passage. Stores all existent passages
|
||||||
|
collection_filepath = os.path.join(data_folder, 'collection.tsv')
|
||||||
|
if not os.path.exists(collection_filepath):
|
||||||
|
tar_filepath = os.path.join(data_folder, 'collection.tar.gz')
|
||||||
|
if not os.path.exists(tar_filepath):
|
||||||
|
logging.info("Download collection.tar.gz")
|
||||||
|
util.http_get('https://msmarco.blob.core.windows.net/msmarcoranking/collection.tar.gz', tar_filepath)
|
||||||
|
|
||||||
|
with tarfile.open(tar_filepath, "r:gz") as tar:
|
||||||
|
tar.extractall(path=data_folder)
|
||||||
|
|
||||||
|
logging.info("Read corpus: collection.tsv")
|
||||||
|
with open(collection_filepath, 'r', encoding='utf8') as fIn:
|
||||||
|
for line in fIn:
|
||||||
|
pid, passage = line.strip().split("\t")
|
||||||
|
pid = int(pid)
|
||||||
|
corpus[pid] = passage
|
||||||
|
|
||||||
|
|
||||||
|
### Read the train queries, store in queries dict
|
||||||
|
queries = {} #dict in the format: query_id -> query. Stores all training queries
|
||||||
|
queries_filepath = os.path.join(data_folder, 'queries.train.tsv')
|
||||||
|
if not os.path.exists(queries_filepath):
|
||||||
|
tar_filepath = os.path.join(data_folder, 'queries.tar.gz')
|
||||||
|
if not os.path.exists(tar_filepath):
|
||||||
|
logging.info("Download queries.tar.gz")
|
||||||
|
util.http_get('https://msmarco.blob.core.windows.net/msmarcoranking/queries.tar.gz', tar_filepath)
|
||||||
|
|
||||||
|
with tarfile.open(tar_filepath, "r:gz") as tar:
|
||||||
|
tar.extractall(path=data_folder)
|
||||||
|
|
||||||
|
|
||||||
|
with open(queries_filepath, 'r', encoding='utf8') as fIn:
|
||||||
|
for line in fIn:
|
||||||
|
qid, query = line.strip().split("\t")
|
||||||
|
qid = int(qid)
|
||||||
|
queries[qid] = query
|
||||||
|
|
||||||
|
|
||||||
|
# Load a dict (qid, pid) -> ce_score that maps query-ids (qid) and paragraph-ids (pid)
|
||||||
|
# to the CrossEncoder score computed by the cross-encoder/ms-marco-MiniLM-L-6-v2 model
|
||||||
|
ce_scores_file = os.path.join(data_folder, 'cross-encoder-ms-marco-MiniLM-L-6-v2-scores.pkl.gz')
|
||||||
|
if not os.path.exists(ce_scores_file):
|
||||||
|
logging.info("Download cross-encoder scores file")
|
||||||
|
util.http_get('https://huggingface.co/datasets/sentence-transformers/msmarco-hard-negatives/resolve/main/cross-encoder-ms-marco-MiniLM-L-6-v2-scores.pkl.gz', ce_scores_file)
|
||||||
|
|
||||||
|
logging.info("Load CrossEncoder scores dict")
|
||||||
|
with gzip.open(ce_scores_file, 'rb') as fIn:
|
||||||
|
ce_scores = pickle.load(fIn)
|
||||||
|
|
||||||
|
# As training data we use hard-negatives that have been mined using various systems
|
||||||
|
hard_negatives_filepath = os.path.join(data_folder, 'msmarco-hard-negatives.jsonl.gz')
|
||||||
|
if not os.path.exists(hard_negatives_filepath):
|
||||||
|
logging.info("Download cross-encoder scores file")
|
||||||
|
util.http_get('https://huggingface.co/datasets/sentence-transformers/msmarco-hard-negatives/resolve/main/msmarco-hard-negatives.jsonl.gz', hard_negatives_filepath)
|
||||||
|
|
||||||
|
|
||||||
|
logging.info("Read hard negatives train file")
|
||||||
|
train_queries = {}
|
||||||
|
negs_to_use = None
|
||||||
|
with gzip.open(hard_negatives_filepath, 'rt') as fIn:
|
||||||
|
for line in tqdm.tqdm(fIn):
|
||||||
|
if max_passages > 0 and len(train_queries) >= max_passages:
|
||||||
|
break
|
||||||
|
data = json.loads(line)
|
||||||
|
|
||||||
|
#Get the positive passage ids
|
||||||
|
pos_pids = data['pos']
|
||||||
|
|
||||||
|
#Get the hard negatives
|
||||||
|
neg_pids = set()
|
||||||
|
if negs_to_use is None:
|
||||||
|
if args.negs_to_use is not None: #Use specific system for negatives
|
||||||
|
negs_to_use = args.negs_to_use.split(",")
|
||||||
|
else: #Use all systems
|
||||||
|
negs_to_use = list(data['neg'].keys())
|
||||||
|
logging.info("Using negatives from the following systems: {}".format(", ".join(negs_to_use)))
|
||||||
|
|
||||||
|
for system_name in negs_to_use:
|
||||||
|
if system_name not in data['neg']:
|
||||||
|
continue
|
||||||
|
|
||||||
|
system_negs = data['neg'][system_name]
|
||||||
|
negs_added = 0
|
||||||
|
for pid in system_negs:
|
||||||
|
if pid not in neg_pids:
|
||||||
|
neg_pids.add(pid)
|
||||||
|
negs_added += 1
|
||||||
|
if negs_added >= num_negs_per_system:
|
||||||
|
break
|
||||||
|
|
||||||
|
if args.use_all_queries or (len(pos_pids) > 0 and len(neg_pids) > 0):
|
||||||
|
train_queries[data['qid']] = {'qid': data['qid'], 'query': queries[data['qid']], 'pos': pos_pids, 'neg': neg_pids}
|
||||||
|
|
||||||
|
logging.info("Train queries: {}".format(len(train_queries)))
|
||||||
|
|
||||||
|
# We create a custom MSMARCO dataset that returns triplets (query, positive, negative)
|
||||||
|
# on-the-fly based on the information from the mined-hard-negatives jsonl file.
|
||||||
|
class MSMARCODataset(Dataset):
|
||||||
|
def __init__(self, queries, corpus, ce_scores):
|
||||||
|
self.queries = queries
|
||||||
|
self.queries_ids = list(queries.keys())
|
||||||
|
self.corpus = corpus
|
||||||
|
self.ce_scores = ce_scores
|
||||||
|
|
||||||
|
for qid in self.queries:
|
||||||
|
self.queries[qid]['pos'] = list(self.queries[qid]['pos'])
|
||||||
|
self.queries[qid]['neg'] = list(self.queries[qid]['neg'])
|
||||||
|
random.shuffle(self.queries[qid]['neg'])
|
||||||
|
|
||||||
|
def __getitem__(self, item):
|
||||||
|
query = self.queries[self.queries_ids[item]]
|
||||||
|
query_text = query['query']
|
||||||
|
qid = query['qid']
|
||||||
|
|
||||||
|
if len(query['pos']) > 0:
|
||||||
|
pos_id = query['pos'].pop(0) #Pop positive and add at end
|
||||||
|
pos_text = self.corpus[pos_id]
|
||||||
|
query['pos'].append(pos_id)
|
||||||
|
else: #We only have negatives, use two negs
|
||||||
|
pos_id = query['neg'].pop(0) #Pop negative and add at end
|
||||||
|
pos_text = self.corpus[pos_id]
|
||||||
|
query['neg'].append(pos_id)
|
||||||
|
|
||||||
|
#Get a negative passage
|
||||||
|
neg_id = query['neg'].pop(0) #Pop negative and add at end
|
||||||
|
neg_text = self.corpus[neg_id]
|
||||||
|
query['neg'].append(neg_id)
|
||||||
|
|
||||||
|
pos_score = self.ce_scores[qid][pos_id]
|
||||||
|
neg_score = self.ce_scores[qid][neg_id]
|
||||||
|
|
||||||
|
return InputExample(texts=[query_text, pos_text, neg_text], label=pos_score-neg_score)
|
||||||
|
|
||||||
|
def __len__(self):
|
||||||
|
return len(self.queries)
|
||||||
|
|
||||||
|
# For training the SentenceTransformer model, we need a dataset, a dataloader, and a loss used for training.
|
||||||
|
train_dataset = MSMARCODataset(queries=train_queries, corpus=corpus, ce_scores=ce_scores)
|
||||||
|
train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=train_batch_size, drop_last=True)
|
||||||
|
train_loss = losses.MarginMSELoss(model=model)
|
||||||
|
|
||||||
|
# Train the model
|
||||||
|
model.fit(train_objectives=[(train_dataloader, train_loss)],
|
||||||
|
epochs=num_epochs,
|
||||||
|
warmup_steps=args.warmup_steps,
|
||||||
|
use_amp=True,
|
||||||
|
checkpoint_path=model_save_path,
|
||||||
|
checkpoint_save_steps=10000,
|
||||||
|
optimizer_params = {'lr': args.lr},
|
||||||
|
)
|
||||||
|
|
||||||
|
# Train latest model
|
||||||
|
model.save(model_save_path)
|
||||||
|
|
||||||
|
# Script was called via:
|
||||||
|
#python train_dense.py
|
||||||
Reference in New Issue
Block a user