ssmits/Falcon2-5.5B-multilingual-embed-base

Usage

Embeddings version of the base model ssmits/Falcon2-5.5B-multilingual. The 'lm_head' layer of this model has been removed, which means it can be used for embeddings. It will not perform greatly, as it needs to be further fine-tuned, as it is pruned and shown by intfloat/e5-mistral-7b-instruct. Additionaly, in stead of a normalization layer, the hidden layers are followed up by both a classical weight and bias 1-dimensional array of 4096 values. The basic Sentence-Transformers implementation is working correctly. This would imply other more sophisticated embeddings techniques such as adding a custom classification head, will work correctly as well.

Inference (sentence-transformers)

from sentence_transformers import SentenceTransformer
import torch

# 1. Load a pretrained Sentence Transformer model
model = SentenceTransformer("ssmits/Falcon2-5.5B-multilingual-embed-base")  # device = "cpu" when <= 24 GB VRAM

# The sentences to encode
sentences = [
    "The weather is lovely today.",
    "It's so sunny outside!",
    "He drove to the stadium.",
]

# 2. Calculate embeddings by calling model.encode()
embeddings = model.encode(sentences)
print(embeddings.shape)
# (3, 4096)

# 3. Calculate the embedding similarities
# Using torch to compute cosine similarity matrix
similarities = torch.nn.functional.cosine_similarity(embeddings.unsqueeze(0), embeddings.unsqueeze(1), dim=2)
print(similarities)
# tensor([[1.0000, 0.7120, 0.5937],
#         [0.7120, 1.0000, 0.5925],
#         [0.5937, 0.5925, 1.0000]])

Note: In my tests it utilizes more than 24GB (RTX 4090), so an A100 or A6000 would be required for inference.

Inference (HuggingFace Transformers)

Without sentence-transformers, 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.

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('ssmits/Falcon2-5.5B-multilingual-embed-base')
model = AutoModel.from_pretrained('ssmits/Falcon2-5.5B-multilingual-embed-base') # device = "cpu" when <= 24 GB VRAM

# 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)

How to enable Multi-GPU

from transformers import AutoModel
from torch.nn import DataParallel

model = AutoModel.from_pretrained("ssmits/Falcon2-5.5B-multilingual-embed-base")
for module_key, module in model._modules.items():
    model._modules[module_key] = DataParallel(module)