<jupyter_start><jupyter_text>Train<jupyter_code># sample training script from https://github.com/UKPLab/beir/blob/main/examples/retrieval/training/train_sbert.py
import torch
import torch.nn as nn

from sentence_transformers import losses, models, SentenceTransformer

from beir import util, LoggingHandler
from beir.datasets.data_loader import GenericDataLoader
from beir.retrieval.train import TrainRetriever


if LABEL == 'QE':
    # load HF model
    word_embedding_model = models.Transformer(model_name, max_seq_length=350)
    
    # add special tokens
    special_tokens_dict = {'additional_special_tokens': ['[Q]','[E]','[R]']}
    num_added_toks = word_embedding_model.tokenizer.add_special_tokens(special_tokens_dict)
    word_embedding_model.auto_model.resize_token_embeddings(len(word_embedding_model.tokenizer))

    pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension())
    model = SentenceTransformer(modules=[word_embedding_model, pooling_model])

else:
    # pretrained sentence-transformer model
    model = SentenceTransformer(model_name)

retriever = TrainRetriever(model=model, batch_size=BATCH_SIZE)

# send model to the GPUs
if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUs!")
retriever = nn.DataParallel(retriever)
retriever.to('cuda')

#### Prepare training samples
train_samples = retriever.module.load_train(corpus, queries, qrels)
train_dataloader = retriever.module.prepare_train(train_samples, shuffle=True)

#### Training SBERT with cosine-product
if similarity == 'cos_sim':
    train_loss = losses.MultipleNegativesRankingLoss(model=retriever.module.model)
#### training SBERT with dot-product
elif similarity == 'dot':
    train_loss = losses.MultipleNegativesRankingLoss(model=retriever.module.model, similarity_fct=util.dot_score)

#### Prepare dev evaluator
# ir_evaluator = retriever.load_ir_evaluator(dev_corpus, dev_queries, dev_qrels)

#### If no dev set is present from above use dummy evaluator
ir_evaluator = retriever.module.load_dummy_evaluator()

#### Configure Train params
evaluation_steps = 5000
warmup_steps = int(len(train_samples) * NEPOCHS / retriever.module.batch_size * 0.1)

retriever.module.fit(train_objectives=[(train_dataloader, train_loss)], 
                                evaluator=ir_evaluator, 
                                epochs=NEPOCHS,
                                output_path=model_save_path,
                                warmup_steps=warmup_steps,
                                evaluation_steps=evaluation_steps,
                                save_best_model=True,
                                use_amp=True)
print(model_save_path, 'trained.')<jupyter_output>Let's use 2 GPUs!<jupyter_text>Evaluate<jupyter_code>from utils import evaluate_model, model_path

# new_model_name = 'msmarco-distilbert-base-tas-b-cos_sim-WD18_entities_wav2vec2-base-960h-512-512-'
evaluate_model(model_path+new_model_name, 'relations-extra')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.0954122340425532, 'fraction_numerical': 0.007978723404255319, 'mean_word_length': 4.373214285714286, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 9, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '361721', 'license': 'MIT', 'n_tokens_mistral': 1013, 'n_tokens_neox': 968, 'n_words': 159, 'path': 'notebooks/fine-tune.ipynb', 'repo_name': 'svakulenk0/spoken_qa', 'size': '3008'}