HuggingFace `Trainer` class compatibility

Hi all,

First off, thank you for developing this awesome library! I have a question regarding the compatibility with the Trainer class in HuggingFace.

Background

I am gathering a set of NLP benchmarks to test some models and have already spent significant time developing my code-base with the HuggingFace transformer library. Some of my newer benchmarks require approaches only available out-of-the-box in sentence-transformers; such as training a bi-encoder with CosineSimilarityLoss. For consistency reasons, I would ideally like to stick to the HuggingFace framework for most of my code-base instead of using two separate frameworks.

Question

In light of this, I was wondering whether it is possible to train/fine-tune a SentenceTransformer by using a modified HuggingFace Trainer class (with CosineSimilarityLoss for example)? Something along the lines of:

# define model components
word_embedding_model = models.Transformer('bert-base-uncased', max_seq_length=256)
pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension())
dense_model = models.Dense(in_features=pooling_model.get_sentence_embedding_dimension(),
out_features=256, activation_function=nn.Tanh())
model = SentenceTransformer(modules=[word_embedding_model, pooling_model, dense_model])

# use modified HuggingFace trainer for training
trainer = CustomTrainer(model=model, ...)
trainer.train()

Apologies if this was already answered elsewhere, I tried to search for similar threads but could not find much.

Thank you for your time.

[atreyasha]

@paluchnuggets Unfortunately not as yet 😕

But I had a look at the some of the methods that need to be overridden in HF’s Trainer class. They list 12 methods, but I think the most important method to override to produce a MWE should be compute_loss. So if we could port the CosineSimilarityLoss into compute_loss, we could already have something to look at.

I’ll probably try this in the next days and report back. Feel free to also add comments/ideas 😃

[avinashronanki]

@atreyasha Thank you, I will take a look