Views
No views yet
1from sentence_transformers import SentenceTransformer
2
3# Load the model
4model = SentenceTransformer('McGill-NLP/AfriE5-Large-instruct')
5
6# Define queries and documents
7# IMPORTANT: Queries require a specific instruction prefix.
8# Documents do not strictly need a prefix, but usage should mirror mE5 conventions.
9query_instruction = "Instruct: Retrieve sentences that are semantically consistent with the input.\nQuery: "
10
11queries = [
12 "What are the key features of AfriMTEB?",
13 "Hali ya hewa ikoje leo?" # Swahili: How is the weather today?
14]
15
16documents = [
17 "AfriMTEB is a benchmark for evaluating text embeddings in African languages.",
18 "Leo kuna jua kali sana." # Swahili: Today it is very sunny.
19]
20
21# Add prefix to queries
22formatted_queries = [query_instruction + q for q in queries]
23
24# Encode
25query_embeddings = model.encode(formatted_queries, normalize_embeddings=True)
26doc_embeddings = model.encode(documents, normalize_embeddings=True)
27
28# Compute similarity
29scores = (query_embeddings @ doc_embeddings.T) * 100
30print(scores)1import torch
2import torch.nn.functional as F
3from transformers import AutoTokenizer, AutoModel
4
5def average_pool(last_hidden_states, attention_mask):
6 last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
7 return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
8
9# Load model and tokenizer
10tokenizer = AutoTokenizer.from_pretrained('McGill-NLP/AfriE5-Large-instruct')
11model = AutoModel.from_pretrained('McGill-NLP/AfriE5-Large-instruct')
12
13# Define input texts
14query_instruction = "Instruct: Retrieve sentences that are semantically consistent with the input.\nQuery: "
15input_texts = [
16 query_instruction + "What is the capital of Nigeria?",
17 "Abuja is the capital city of Nigeria.",
18 "Lagos is the largest city in Nigeria."
19]
20
21# Tokenize
22batch_dict = tokenizer(input_texts, max_length=512, padding=True, truncation=True, return_tensors='pt')
23
24# Get embeddings
25outputs = model(**batch_dict)
26embeddings = average_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
27
28# Normalize embeddings
29embeddings = F.normalize(embeddings, p=2, dim=1)
30
31# Compute cosine similarity
32scores = (embeddings[:1] @ embeddings[1:].T) * 100
33print(scores)| Model | Average Score |
|---|---|
| AfriE5-Large-instruct | 63.7 |
| Gemini Embedding-001 | 63.1 |
| mE5-Large-instruct | 62.0 |
| BGE-M3 | 55.0 |
| Model | Average Score |
|---|---|
| AfriE5-Large-instruct | 62.4 |
| mE5-Large-instruct | 61.3 |
| Gemini Embedding-001 | 60.6 |
| BGE-M3 | 55.8 |
facebook/nllb-200-3.3B.BAAI/bge-reranker-v2-m3.1@article{uemura2025afrimteb,
2 title={AfriMTEB and AfriE5: Benchmarking and Adapting Text Embedding Models for African Languages},
3 author={Uemura, Kosei and Zhang, Miaoran and Adelani, David Ifeoluwa},
4 journal={arXiv preprint},
5 year={2025}
6}