ArmEmbed - Text Embedding Model for Armenian

This embedding model is built on a Llama-based language model pre-trained on Armenian text and further adapted using LoRA on additional Armenian data. It produces 2048-dimensional embeddings.

Usage

Below is an example to encode queries and passages from the MS-MARCO passage ranking dataset.

Sentence Transformers

from sentence_transformers import SentenceTransformer

queries = [
    "Ի՞նչ է պատահում հյուսվածքներին, ինչպիսիք են սիրտը կամ ուղեղը, եթե թթվածնով հարուստ արյունը ժամանակին չի մատակարարվում:",
    "Կարո՞ղ է արդյոք ոստիկանությունը հետախուզման թույլտվություն ստանալ այն բանից հետո, երբ նրանք տեսել են ապացույցներ:",
]

passages = [
    "Եվ․․․ Կիսալուսնաձև փականները կանխում են արյան հետհոսքը զարկերակներից դեպի փորոքներ։ Բացատրեք, թե ինչ է պատահում հյուսվածքներին, ինչպիսիք են սիրտը կամ ուղեղը, եթե թթվածնով հարուստ արյունը ժամանակին չի մատակարարվում: Օրգանի հյուսվածքը սկսում է մահանալ։",
    "Ոստիկանությունը կտրամադրի իր սեփական ապացույցները հետախուզման թույլտվության համար, և կասկածյալը ներկա չէ, երբ թույլտվություն է տրվում: Երբ հետախուզման թույլտվություն է ստացվում, ոստիկանությունը կարող է խուզարկել միայն թույլտվության մեջ նշված վայրը, լինի դա տուն, մեքենա, թե որոշակի արտաքին վայր:",
]

prefixed_queries = ["query: " + query for query in queries]
prefixed_passages = [" passage: " + passage for passage in passages]

sentence_transformer = SentenceTransformer("alexshah/armembed", trust_remote_code=True)

query_embeddings = sentence_transformer.encode(
    prefixed_queries, normalize_embeddings=True
)
passage_embeddings = sentence_transformer.encode(
    prefixed_passages, normalize_embeddings=True
)

scores = (query_embeddings @ passage_embeddings.T) * 100
print(scores.tolist())

Transformers

import torch
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer


def last_token_pool(last_hidden_states, attention_mask):
    left_padding = attention_mask[:, -1].sum() == attention_mask.shape[0]
    if left_padding:
        return last_hidden_states[:, -1]
    else:
        sequence_lengths = attention_mask.sum(dim=1) - 1
        batch_size = last_hidden_states.shape[0]
        return last_hidden_states[
            torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths
        ]


queries = [
    "Ի՞նչ է պատահում հյուսվածքներին, ինչպիսիք են սիրտը կամ ուղեղը, եթե թթվածնով հարուստ արյունը ժամանակին չի մատակարարվում:",
    "Կարո՞ղ է արդյոք ոստիկանությունը հետախուզման թույլտվություն ստանալ այն բանից հետո, երբ նրանք տեսել են ապացույցներ:",
]

passages = [
    "Եվ․․․ Կիսալուսնաձև փականները կանխում են արյան հետհոսքը զարկերակներից դեպի փորոքներ։ Բացատրեք, թե ինչ է պատահում հյուսվածքներին, ինչպիսիք են սիրտը կամ ուղեղը, եթե թթվածնով հարուստ արյունը ժամանակին չի մատակարարվում: Օրգանի հյուսվածքը սկսում է մահանալ։",
    "Ոստիկանությունը կտրամադրի իր սեփական ապացույցները հետախուզման թույլտվության համար, և կասկածյալը ներկա չէ, երբ թույլտվություն է տրվում: Երբ հետախուզման թույլտվություն է ստացվում, ոստիկանությունը կարող է խուզարկել միայն թույլտվության մեջ նշված վայրը, լինի դա տուն, մեքենա, թե որոշակի արտաքին վայր:",
]

input_texts = ["query: " + query for query in queries] + [
    " passage: " + passage for passage in passages
]

model = AutoModel.from_pretrained("alexshah/armembed")
tokenizer = AutoTokenizer.from_pretrained("alexshah/armembed")

batch_dict = tokenizer(
    input_texts, max_length=512, padding=True, truncation=True, return_tensors="pt"
)
outputs = model(**batch_dict)
embeddings = last_token_pool(outputs.last_hidden_state, batch_dict["attention_mask"])

embeddings = F.normalize(embeddings, p=2, dim=1)
scores = (embeddings[:2] @ embeddings[2:].T) * 100
print(scores.tolist())

Intended Use

Primary Intended Uses

• Retrieval-augmented generation (RAG)
• Semantic search in Armenian
• Document similarity computation
• Cross-lingual text understanding
• Text classification tasks
• Information retrieval

Training Data

Dataset Details

• Source: Reddit dataset with English-Armenian translations
• Size: 0.66M pairs of rows
• Content Type: Title and body text pairs
• Split Ratio: 98.5% train, 1.5% test

Training Procedure

Training Details

• LoRA: 16/32/0.1 on all linear layers
• Training Duration: ~15 hours
• Hardware: 4 × NVIDIA A100 (40GB) GPUs
• Effective Batch Size: 128 (2 per GPU × 4 GPUs, with gradient accumulation of 16)
• Learning Rate: 1e-4
• Weight Decay: 0.01
• Warmup Steps: 200
• Max Sequence Length: 512 tokens
• FP16 Training: Enabled
• Gradient Clipping: 1.0

Optimization Configuration

• Framework: DeepSpeed ZeRO Stage 3
• Optimizer: AdamW with auto learning rate and weight decay
• Mixed Precision: bfloat16 (bf16) enabled
• ZeRO Optimization: Stage 3 with:
  • No parameter offloading
  • Overlap communication enabled
  • Contiguous gradients enabled
  • Auto-tuned reduce and prefetch bucket sizes
  • Auto-tuned parameter persistence threshold

Technical Specifications

• Model Size: ~1.24B parameters (based on a Llama-style architecture LLM pre-trained on Armenian data)
• Embedding Dimension: 2048
• Max Sequence Length: 512 tokens
• Framework Compatibility:
  • PyTorch
  • Hugging Face Transformers
  • DeepSpeed

Model Details

• Model Name: ArmEmbed
• Model Type: Text Embeddings for the Armenian Language
• Version: 1.0.0
• License: Apache 2.0
• Last Updated: May 2025
• Model Architecture: Transformer-based embeddings model
• Input: Armenian text
• Output: Dense vector embeddings (size=2048)