初始化项目，由ModelHub XC社区提供模型

Model: sentence-transformers/msmarco-distilbert-dot-v5
Source: Original Platform

.gitattributes

@@ -0,0 +1,28 @@
+*.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
+model.safetensors filter=lfs diff=lfs merge=lfs -text

1_Pooling/config.json

@@ -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
+}

README.md

@@ -0,0 +1,190 @@
+---
+language:
+- en
+license: apache-2.0
+library_name: sentence-transformers
+tags:
+- sentence-transformers
+- feature-extraction
+- sentence-similarity
+- transformers
+pipeline_tag: sentence-similarity
+---
+
+# msmarco-distilbert-dot-v5
+This is a [sentence-transformers](https://www.SBERT.net) model: It maps sentences & paragraphs to a 768 dimensional dense vector space and was designed for **semantic search**. It has been trained on 500K (query, answer) pairs from the [MS MARCO dataset](https://github.com/microsoft/MSMARCO-Passage-Ranking/). For an introduction to semantic search, have a look at: [SBERT.net - Semantic Search](https://www.sbert.net/examples/applications/semantic-search/README.html)
+
+
+## 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, util
+
+query = "How many people live in London?"
+docs = ["Around 9 Million people live in London", "London is known for its financial district"]
+
+#Load the model
+model = SentenceTransformer('sentence-transformers/msmarco-distilbert-dot-v5')
+
+#Encode query and documents
+query_emb = model.encode(query)
+doc_emb = model.encode(docs)
+
+#Compute dot score between query and all document embeddings
+scores = util.dot_score(query_emb, doc_emb)[0].cpu().tolist()
+
+#Combine docs & scores
+doc_score_pairs = list(zip(docs, scores))
+
+#Sort by decreasing score
+doc_score_pairs = sorted(doc_score_pairs, key=lambda x: x[1], reverse=True)
+
+#Output passages & scores
+print("Query:", query)
+for doc, score in doc_score_pairs:
+    print(score, doc)
+```
+
+
+## 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 correct 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.last_hidden_state
+    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)
+
+
+#Encode text
+def encode(texts):
+    # Tokenize sentences
+    encoded_input = tokenizer(texts, padding=True, truncation=True, return_tensors='pt')
+
+    # Compute token embeddings
+    with torch.no_grad():
+        model_output = model(**encoded_input, return_dict=True)
+
+    # Perform pooling
+    embeddings = mean_pooling(model_output, encoded_input['attention_mask'])
+
+    return embeddings
+
+
+# Sentences we want sentence embeddings for
+query = "How many people live in London?"
+docs = ["Around 9 Million people live in London", "London is known for its financial district"]
+
+# Load model from HuggingFace Hub
+tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/msmarco-distilbert-dot-v5")
+model = AutoModel.from_pretrained("sentence-transformers/msmarco-distilbert-dot-v5")
+
+#Encode query and docs
+query_emb = encode(query)
+doc_emb = encode(docs)
+
+#Compute dot score between query and all document embeddings
+scores = torch.mm(query_emb, doc_emb.transpose(0, 1))[0].cpu().tolist()
+
+#Combine docs & scores
+doc_score_pairs = list(zip(docs, scores))
+
+#Sort by decreasing score
+doc_score_pairs = sorted(doc_score_pairs, key=lambda x: x[1], reverse=True)
+
+#Output passages & scores
+print("Query:", query)
+for doc, score in doc_score_pairs:
+    print(score, doc)
+```
+
+## Technical Details
+
+In the following some technical details how this model must be used:
+
+| Setting | Value |
+| --- | :---: |
+| Dimensions | 768 |
+| Max Sequence Length | 512 |
+| Produces normalized embeddings | No |
+| Pooling-Method | Mean pooling |
+| Suitable score functions | dot-product (e.g. `util.dot_score`) |
+
+
+## Training
+
+See `train_script.py` in this repository for the used training script.
+
+
+
+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:
+```
+{
+    "callback": null,
+    "epochs": 30,
+    "evaluation_steps": 0,
+    "evaluator": "NoneType",
+    "max_grad_norm": 1,
+    "optimizer_class": "<class 'transformers.optimization.AdamW'>",
+    "optimizer_params": {
+        "lr": 1e-05
+    },
+    "scheduler": "WarmupLinear",
+    "steps_per_epoch": null,
+    "warmup_steps": 10000,
+    "weight_decay": 0.01
+}
+```
+
+
+## Full Model Architecture
+```
+SentenceTransformer(
+  (0): Transformer({'max_seq_length': 512, '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
+
+This model was trained by [sentence-transformers](https://www.sbert.net/). 
+        
+If you find this model helpful, feel free to cite our publication [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://arxiv.org/abs/1908.10084):
+```bibtex 
+@inproceedings{reimers-2019-sentence-bert,
+    title = "Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks",
+    author = "Reimers, Nils and Gurevych, Iryna",
+    booktitle = "Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing",
+    month = "11",
+    year = "2019",
+    publisher = "Association for Computational Linguistics",
+    url = "http://arxiv.org/abs/1908.10084",
+}
+```
+
+## License
+
+This model is released under the Apache 2 license. However, note that this model was trained on the MS MARCO dataset which has it's own license restrictions: [MS MARCO - Terms and Conditions](https://github.com/microsoft/msmarco/blob/095515e8e28b756a62fcca7fcf1d8b3d9fbb96a9/README.md).

config.json

@@ -0,0 +1,23 @@
+{
+  "_name_or_path": "final-models/distilbert-margin_mse-sym_mnrl-mean-v1/",
+  "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,
+  "transformers_version": "4.6.1",
+  "vocab_size": 30522
+}

config_sentence_transformers.json

@@ -0,0 +1,10 @@
+{
+  "__version__": {
+    "sentence_transformers": "3.0.0.dev0",
+    "transformers": "4.41.0.dev0",
+    "pytorch": "2.3.0+cu121"
+  },
+  "prompts": {},
+  "default_prompt_name": null,
+  "similarity_fn_name": "dot"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7daa9b9afe07570758bb810eb55b3f248a672c7c348396907d7ce46630cff0df
+size 265462608

modules.json

@@ -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"
+  }
+]

onnx/model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e9ceed39618b9c8bcd3d282fe42f5fb285ba7f9ac3f0b8d365db97b615eafe6e
+size 265564612

onnx/model_O1.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9f731d779bfcf24e1d969b312c5b888c6ad96f08574a0715801721a79795f76b
+size 265548916

onnx/model_O2.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:359569060665cbcab0b79b62433765e05375a6fb71b61549f11300c9312d46b7
+size 265471026

onnx/model_O3.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b1d2aabab430cee551b57d33e1e13316b92df9dd3f70c5e1e47f2eb7d2f06a0f
+size 265471086

onnx/model_O4.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ef1ffa77578c058aea1a09be4471698f739f17b4d639eea831084d632ba58ccf
+size 132748309

onnx/model_qint8_arm64.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8c86d6e7935918409b7be0ed2e4e16a4ad026cf621894870d678c6a984192e7a
+size 66965107

onnx/model_qint8_avx512.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8c86d6e7935918409b7be0ed2e4e16a4ad026cf621894870d678c6a984192e7a
+size 66965107

onnx/model_qint8_avx512_vnni.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8c86d6e7935918409b7be0ed2e4e16a4ad026cf621894870d678c6a984192e7a
+size 66965107

onnx/model_quint8_avx2.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f6a236a06b0c2fa3db0a8bc1e0dca79f12752c60b20e87edcd12a3f9e6e35059
+size 67006579

openvino/openvino_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:791914376291e0f2309b3d0ec5b37a02aef8b89ed6a7f041810d62fd3d0ceb10
+size 265455736

openvino/openvino_model_qint8_quantized.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:591fdf409e3a061373a096ae90986716b490e57a47a0fccba82df60c89e4010e
+size 66970752

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6abdf990eaa861287cecd295bc6dfa69d8fdc89007ad49b1179de73412e46e05
+size 265491187

sentence_bert_config.json

@@ -0,0 +1,4 @@
+{
+  "max_seq_length": 512,
+  "do_lower_case": false
+}

special_tokens_map.json

@@ -0,0 +1 @@
+{"unk_token": "[UNK]", "sep_token": "[SEP]", "pad_token": "[PAD]", "cls_token": "[CLS]", "mask_token": "[MASK]"}

tf_model.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:730a6b8f936ed95bd2c38fbe899a3816102357a32cc2e1a6aff38ee93e3d6a65
+size 265571968

tokenizer_config.json

@@ -0,0 +1 @@
+{"do_lower_case": true, "unk_token": "[UNK]", "sep_token": "[SEP]", "pad_token": "[PAD]", "cls_token": "[CLS]", "mask_token": "[MASK]", "tokenize_chinese_chars": true, "strip_accents": null, "model_max_length": 512, "special_tokens_map_file": null, "name_or_path": "final-models/distilbert-margin_mse-sym_mnrl-mean-v1/"}

train_script.py

@@ -0,0 +1,239 @@
+
+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
+from collections import defaultdict
+from torch.utils.data import IterableDataset
+import tqdm
+from torch.utils.data import Dataset
+import random
+from shutil import copyfile
+
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--train_batch_size", default=64, type=int)
+parser.add_argument("--max_seq_length", default=300, type=int)
+parser.add_argument("--model_name", required=True)
+parser.add_argument("--max_passages", default=0, type=int)
+parser.add_argument("--epochs", default=10, type=int)
+parser.add_argument("--pooling", default="cls")
+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("--name", default='')
+parser.add_argument("--num_negs_per_system", default=5, type=int) 
+parser.add_argument("--use_pre_trained_model", default=False, action="store_true") 
+parser.add_argument("--use_all_queries", default=False, action="store_true") 
+args = parser.parse_args()
+
+print(args)
+
+#### 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
+
+# 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
+
+# We construct the SentenceTransformer bi-encoder from scratch
+if args.use_pre_trained_model:
+    print("use pretrained SBERT model")
+    model = SentenceTransformer(model_name)
+    model.max_seq_length = max_seq_length
+else:
+    print("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])
+
+model_save_path = f'output/train_bi-encoder-margin_mse_en-{args.name}-{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")
+        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")
+        queries[qid] = query
+
+
+# Read our training file: msmarco-hard-negatives.jsonl.gz contains all queries and hard-negatives that were mined with different systems
+# For each positive and mined-hard negative passage, we have a Cross-Encoder score from the cross-encoder/ms-marco-MiniLM-L-6-v2 model
+# This  Cross-Encoder score allows to de-noise our hard-negatives by requiring that their CE-score is below a certain treshold
+train_filepath = '/home/msmarco/data/hard-negatives/msmarco-hard-negatives-v6.jsonl.gz'
+
+#### Create our training data
+logging.info("Read train dataset")
+train_queries = {}
+ce_scores = {}
+negs_to_use = None
+with gzip.open(train_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)
+        
+        if data['qid'] not in ce_scores:
+            ce_scores[data['qid']] = {}
+        
+        # Add pos ce_scores
+        for item in data['pos'] :
+            ce_scores[data['qid']][item['pid']] = item['ce-score']
+
+        #Get the positive passage ids
+        pos_pids = [item['pid'] for item in 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())
+            print("Using negatives from the following systems:", 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 item in system_negs:
+                #Add neg ce_scores
+                ce_scores[data['qid']][item['pid']] = item['ce-score']
+                
+                pid = item['pid']
+                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,
+          checkpoint_save_total_limit = 0,
+          optimizer_params = {'lr': args.lr},
+          )
+
+# Train latest model
+model.save(model_save_path)
+
+
+# Script was called via:
+#python train_bi-encoder-margin_mse-en.py --model final-models/distilbert-margin_mse-sym_mnrl-mean-v1 --lr=1e-5 --warmup_steps=10000 --negs_to_use=distilbert-margin_mse-sym_mnrl-mean-v1 --num_negs_per_system=10 --epochs=30 --name=cnt_with_mined_negs_mean --use_pre_trained_model --train_batch_size 64