1from matryoshka_wrapper import load_model, MatryoshkaWrapper
2from torch.nn.functional import cosine_similarity
3
4# Load model with desired dimension
5repo_name = "Abdalrahmankamel/matryoshka-arabert"
6model, tokenizer = load_model(repo_name, dim="256")
7
8# Example texts
9text1 = "هذا المنتج كان مخيبًا للآمال."
10text2 = "هذه البضاعة رائعة!"
11
12# Generate embeddings
13emb1 = model.get_embedding(text1, tokenizer, dim="256").squeeze()
14emb2 = model.get_embedding(text2, tokenizer, dim="256").squeeze()
15
16# Calculate similarity
17similarity = cosine_similarity(emb1.unsqueeze(0), emb2.unsqueeze(0)).item()
18print(f"🔍 Cosine Similarity: {similarity:.4f}")
1# Triplet data example
2anchor = "الطفل يلعب في الحديقة"
3positive = "ولد صغير يلهو في البستان" # Similar sentence
4negative = "السيارة تسير في الشارع" # Different sentence
5
6# Generate embeddings
7anchor_emb = model.get_embedding(anchor, tokenizer, dim="256").squeeze()
8positive_emb = model.get_embedding(positive, tokenizer, dim="256").squeeze()
9negative_emb = model.get_embedding(negative, tokenizer, dim="256").squeeze()
10
11# Calculate similarities
12sim_positive = cosine_similarity(anchor_emb.unsqueeze(0), positive_emb.unsqueeze(0)).item()
13sim_negative = cosine_similarity(anchor_emb.unsqueeze(0), negative_emb.unsqueeze(0)).item()
14
15print("🔍 Triplet Results:")
16print(f"📊 Anchor ↔ Positive: {sim_positive:.4f}")
17print(f"📊 Anchor ↔ Negative: {sim_negative:.4f}")
18print(f"📈 Margin: {sim_positive - sim_negative:.4f}")
19
20if sim_positive > sim_negative:
21 print("✅ Triplet Success!")
22else:
23 print("❌ Triplet Failed!")
1# Compare across all dimensions
2text1 = "اطفال يمرحون سوياً بالكرة في المساحات الخضراء"
3text2 = "أطفال يلعبون كرة القدم على العشب"
4
5dimensions = [8, 64, 128, 256]
6
7print("🔍 Multi-Dimensional Similarity Comparison")
8print("=" * 50)
9
10for dim in dimensions:
11 model, tokenizer = load_model(repo_name, dim=str(dim))
12
13 emb1 = model.get_embedding(text1, tokenizer, dim=str(dim)).squeeze()
14 emb2 = model.get_embedding(text2, tokenizer, dim=str(dim)).squeeze()
15
16 similarity = cosine_similarity(emb1.unsqueeze(0), emb2.unsqueeze(0)).item()
17
18 print(f"📐 Dim {dim:>3}: Similarity = {similarity:.4f} | Shape = {emb1.shape}")
1@misc{kamel2025arabert-matryoshka,
2 author = {Abdalrahman Kamel},
3 title = {AraBERT Matryoshka: Multi-Dimensional Arabic Sentence Embeddings with Triplet Loss},
4 year = {2025},
5 url = {https://huggingface.co/Abdalrahmankamel/matryoshka-arabert},
6 note = {Hugging Face Model Repository}
7}
