This is a port of the multilingual SONAR text encoder (
https://huggingface.co/facebook/SONAR) to the
transformers format from
fairseq2.
Its embeddings are expected be equal to those the official implementation (
https://github.com/facebookresearch/SONAR), but the latter stays the source of truth.
The encoder supports the same 202 languages as
NLLB-200
(see also
the source model card
and
FLORES-200 lang code mapping).
1# !pip install transformers sentencepiece -q
2
3import torch
4from transformers import AutoTokenizer
5from transformers.models.m2m_100.modeling_m2m_100 import M2M100Encoder
6
7model_name = "cointegrated/SONAR_200_text_encoder"
8encoder = M2M100Encoder.from_pretrained(model_name)
9tokenizer = AutoTokenizer.from_pretrained(model_name)
10
11def encode_mean_pool(texts, tokenizer, encoder, lang='eng_Latn', norm=False):
12 tokenizer.src_lang = lang
13 with torch.inference_mode():
14 batch = tokenizer(texts, return_tensors='pt', padding=True)
15 seq_embs = encoder(**batch).last_hidden_state
16 mask = batch.attention_mask
17 mean_emb = (seq_embs * mask.unsqueeze(-1)).sum(1) / mask.unsqueeze(-1).sum(1)
18 if norm:
19 mean_emb = torch.nn.functional.normalize(mean_emb)
20 return mean_emb
21
22sentences = ['My name is SONAR.', 'I can embed the sentences into vectorial space.']
23embs = encode_mean_pool(sentences, tokenizer, encoder, lang="eng_Latn")
24print(embs.shape)
25# torch.Size([2, 1024])
26print(embs)
27# tensor([[-0.0053, 0.0020, -0.0006, ..., 0.0094, -0.0009, 0.0070],
28# [-0.0003, -0.0071, 0.0076, ..., 0.0055, 0.0022, -0.0083]])
For advanced examples of usage, please take a look at the readme in
https://github.com/facebookresearch/SONAR.