|
|
--- |
|
|
library_name: transformers |
|
|
language: |
|
|
- hy |
|
|
tags: |
|
|
- sentence-transformers |
|
|
--- |
|
|
|
|
|
# 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 |
|
|
```python |
|
|
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 |
|
|
```python |
|
|
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) |
