🚂 Training and Finetuning

Train powerful sentence embedding models with AnglE using either CLI or Python API.

---

🗂️ Data Preparation

AnglE supports three dataset formats. Choose based on your task:

Format A: Pair with Label

A pair format with three columns: text1, text2, and label. The label should be a similarity score (e.g., 0-1).

Example:

{"text1": "A plane is taking off.", "text2": "An air plane is taking off.", "label": 0.95}

Format B: Query-Positive

A pair format with two columns: query and positive. Both fields can be str or List[str] (random sampling for lists).

Example:

{"query": "A person on a horse jumps over a broken down airplane.", "positive": "A person is outdoors, on a horse."}

Format C: Query-Positive-Negative

A triple format with three columns: query, positive, and negative. All fields can be str or List[str] (random sampling for lists).

Example:

{"query": "Two blond women are hugging one another.", "positive": "There are women showing affection.", "negative": "Men are fighting."}

Note

All formats use HuggingFace datasets.Dataset.

---

🎯 Training Methods

⭐ Method 1: CLI Training (Recommended)

Use angle-trainer to train your models with a simple command-line interface.

Single GPU Training:

CUDA_VISIBLE_DEVICES=0 angle-trainer --help

Multi-GPU Training:

CUDA_VISIBLE_DEVICES=0,1,2,3 WANDB_MODE=disabled accelerate launch \
--multi_gpu \
--num_processes 4 \
--main_process_port 2345 \
-m angle_emb.angle_trainer --help

use FSDP for bigger batch size:

CUDA_VISIBLE_DEVICES=0,1,2,3 WANDB_MODE=disabled accelerate launch \
--multi_gpu \
--num_processes 4 \
--main_process_port 2345 \
--config_file examples/FSDP/fsdp_config.yaml \
-m angle_emb.angle_trainer \
--gradient_checkpointing 1 \
--use_reentrant 0 \
...

see more examples in examples/FSDP

📝 Training Examples

Example 1: BERT-based Model

Train a BERT model with multi-GPU support:

CUDA_VISIBLE_DEVICES=0,1,2,3 WANDB_MODE=disabled accelerate launch \
--multi_gpu \
--num_processes 4 \
--main_process_port 2345 \
-m angle_emb.angle_trainer \
--model_name_or_path WhereIsAI/UAE-Large-V1 \
--train_name_or_path SeanLee97/nli_for_simcse \
--save_dir ckpts/uae-nli \
--column_rename_mapping "text:query" \
--query_prompt "query: {text}" \
--doc_prompt "doc: {text}" \
--learning_rate 1e-5 \
--pooling_strategy cls \
--epochs 1 \
--batch_size 32 \
--logging_steps 10 \
--gradient_accumulation_steps 2 \
--ibn_w 1.0 \
--cln_w 1.0 \
--angle_w 0.02 \
--fp16 1

Example 2: ModernBERT-based Model

Train with ModernBERT architecture:

CUDA_VISIBLE_DEVICES=0,1,2,3 WANDB_MODE=disabled accelerate launch \
--multi_gpu \
--num_processes 4 \
--main_process_port 2345 \
-m angle_emb.angle_trainer \
--model_name_or_path answerdotai/ModernBERT-base \
--train_name_or_path SeanLee97/nli_for_simcse \
--save_dir ckpts/modernbert-nli \
--column_rename_mapping "text:query" \
--query_prompt "query: {text}" \
--doc_prompt "doc: {text}" \
--learning_rate 1e-4 \
--pooling_strategy mean \
--epochs 1 \
--batch_size 128 \
--logging_steps 10 \
--gradient_accumulation_steps 2 \
--ibn_w 1.0 \
--cln_w 1.0 \
--angle_w 0.02 \
--fp16 1

Example 3: LLM-based Model with FSDP

Train large language models using Fully Sharded Data Parallel:

CUDA_VISIBLE_DEVICES=0,1,2,3 WANDB_MODE=disabled accelerate launch \
--multi_gpu \
--num_processes 4 \
--main_process_port 2345 \
--config_file examples/FSDP/fsdp_config.yaml \
-m angle_emb.angle_trainer \
--gradient_checkpointing 1 \
--use_reentrant 0 \
--model_name_or_path Qwen/Qwen3-0.6B \
--torch_dtype "bfloat16" \
--is_llm 1 \
--apply_lora 1 --lora_r 32 --lora_alpha 32 \
--maxlen 312 \
--train_name_or_path SeanLee97/nli_for_simcse \
--save_dir ckpts/qwen-nli \
--column_rename_mapping "text:query" \
--query_prompt "query: {text}" \
--doc_prompt "doc: {text}" \
--learning_rate 1e-4 \
--pooling_strategy last \
--epochs 1 \
--batch_size 16 \
--logging_steps 10 \
--gradient_accumulation_steps 2 \
--ibn_w 1.0 \
--cln_w 1.0 \
--angle_w 0.02 \
--bf16 1

🐍 Method 2: Python API Training

Train models programmatically using the angle_emb library.

Example:

from datasets import load_dataset
from angle_emb import AnglE

# Step 1: Load pretrained model
angle = AnglE.from_pretrained(
    'SeanLee97/angle-bert-base-uncased-nli-en-v1',
    max_length=128,
    pooling_strategy='cls'
).cuda()

# Step 2: Prepare dataset (Format A example)
ds = load_dataset('mteb/stsbenchmark-sts')
ds = ds.map(lambda obj: {
    "text1": str(obj["sentence1"]),
    "text2": str(obj['sentence2']),
    "label": obj['score']
})
ds = ds.select_columns(["text1", "text2", "label"])

# Step 3: Train the model
angle.fit(
    train_ds=ds['train'].shuffle(),
    valid_ds=ds['validation'],
    output_dir='ckpts/sts-b',
    batch_size=32,
    epochs=5,
    learning_rate=2e-5,
    save_steps=100,
    eval_steps=1000,
    warmup_steps=0,
    gradient_accumulation_steps=1,
    loss_kwargs={
        'cosine_w': 1.0,
        'ibn_w': 1.0,
        'angle_w': 0.02,
        'cosine_tau': 20,
        'ibn_tau': 20,
        'angle_tau': 20
    },
    fp16=True,
    logging_steps=100
)

# Step 4: Evaluate
corrcoef = angle.evaluate(ds['test'])
print('Spearman\'s corrcoef:', corrcoef)

---

⚙️ Configuration & Hyperparameters

💡 Loss Weight Parameters

Parameter	Default Value	Description
angle_w	0.02	Weight for angle loss
ibn_w	1.0	Weight for in-batch negative loss
cln_w	1.0	Weight for contrastive learning loss
cosine_w	0.0	Weight for cosine loss

💡 Temperature Parameters

Parameter	Default Value	Description
angle_tau	20.0	Temperature for angle loss
ibn_tau	20.0	Temperature for ibn and cln losses
cosine_tau	20.0	Temperature for cosine loss

💡 Fine-tuning Tips

Format-specific Recommendations:

Prevent Catastrophic Forgetting:

To alleviate information forgetting during fine-tuning:

• Set teacher_name_or_path for knowledge distillation

• Use same model path for self-distillation

• Important: Teacher and student must use the same tokenizer

⚙️ Advanced Features

Training Special Models:

Applying Prompts:

Format	Flag	Applies To
Format A	--text_prompt	Both text1 and text2
Format B/C	--query_prompt	query field
Format B/C	--doc_prompt	positive and negative

Model Conversion:

Convert trained models to sentence-transformers format:

python scripts/convert_to_sentence_transformers.py --help

🔄 Integration with sentence-transformers

Training:

SentenceTransformers provides an AnglE loss implementation.

Warning

The SentenceTransformers implementation is partial. For best results, use the official angle_emb library.

Inference:

Models trained with angle_emb can be converted to sentence-transformers format using the conversion script at examples/convert_to_sentence_transformers.py.

📚 Additional Resources

• Check out the complete 👨‍🏫 Tutorial for a hands-on example

• Learn about 🎯 Evaluation methods

• Explore available 🏛️ Official Pretrained Models

• See 🚀 Quick Start for basic usage