Modernbert Base Long Context Qe V1

🚀 Quality Estimation for Machine Translation

This model provides a solution for reference - free, long - context quality estimation of machine translation systems. It is fine - tuned on a specific dataset, offering reliable evaluation metrics for machine translation quality.

🚀 Quick Start

Installation

1. Install the required libraries.

pip3 install --upgrade datasets accelerate transformers
pip3 install --upgrade flash_attn triton

Inference

2. Load the test dataset.

from datasets import load_dataset

test_dataset = load_dataset("ymoslem/wmt-da-human-evaluation",
                             split="test",
                             trust_remote_code=True
                            )
print(test_dataset)

3. Load the model and tokenizer:

from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

# Load the fine-tuned model and tokenizer
model_name = "ymoslem/ModernBERT-base-long-context-qe-v1"
model = AutoModelForSequenceClassification.from_pretrained(
    model_name,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Move model to GPU if available
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
model.eval()

4. Prepare the dataset. Each source segment src and target segment tgt are separated by the sep_token, which is '</s>' for ModernBERT.

sep_token = tokenizer.sep_token
input_test_texts = [f"{src} {sep_token} {tgt}" for src, tgt in zip(test_dataset["src"], test_dataset["mt"])]

5. Generate predictions.

Basic Usage

If you print model.config.problem_type, the output is regression. Still, you can use the "text - classification" pipeline as follows (cf. pipeline documentation):

from transformers import pipeline

classifier = pipeline("text-classification",
                      model=model_name,
                      tokenizer=tokenizer,
                      device=0,
                     )

predictions = classifier(input_test_texts,
                         batch_size=128,
                         truncation=True,
                         padding="max_length",
                         max_length=tokenizer.model_max_length,
                       )
predictions = [prediction["score"] for prediction in predictions]

Advanced Usage

Alternatively, you can use an elaborate version of the code, which is slightly faster and provides more control.

from torch.utils.data import DataLoader
import torch
from tqdm.auto import tqdm

# Tokenization function
def process_batch(batch, tokenizer, device):
    sep_token = tokenizer.sep_token
    input_texts = [f"{src} {sep_token} {tgt}" for src, tgt in zip(batch["src"], batch["mt"])]
    tokens = tokenizer(input_texts,
                       truncation=True,
                       padding="max_length",
                       max_length=tokenizer.model_max_length,
                       return_tensors="pt",
                      ).to(device)
    return tokens
    


# Create a DataLoader for batching
test_dataloader = DataLoader(test_dataset, 
                             batch_size=128,   # Adjust batch size as needed
                             shuffle=False)


# List to store all predictions
predictions = []

with torch.no_grad():
    for batch in tqdm(test_dataloader, desc="Inference Progress", unit="batch"):

        tokens = process_batch(batch, tokenizer, device)
        
        # Forward pass: Generate model's logits
        outputs = model(**tokens)

        # Get logits (predictions)
        logits = outputs.logits

        # Extract the regression predicted values
        batch_predictions = logits.squeeze()

        # Extend the list with the predictions
        predictions.extend(batch_predictions.tolist())

✨ Features

• Multilingual Support: Supports multiple languages including multilingual, bn, cs, de, en, et, fi, fr, gu, ha, hi, is, ja, kk, km, lt, lv, pl, ps, ru, ta, tr, uk, xh, zh, zu.
• Long - Context Quality Estimation: Trained on long - context datasets, suitable for document - level quality estimation of machine translation systems.

📦 Installation

1. Install the required libraries.

pip3 install --upgrade datasets accelerate transformers
pip3 install --upgrade flash_attn triton

📚 Documentation

Model description

This model is for reference - free, long - context quality estimation (QE) of machine translation (MT) systems. It is trained on a dataset of translation pairs comprising up to 32 sentences (64 sentences for the source and target). Hence, this model is suitable for document - level quality estimation.

Training and evaluation data

The model is trained on the long - context dataset [ymoslem/wmt - da - human - evaluation - long - context](https://huggingface.co/datasets/ymoslem/wmt - da - human - evaluation - long - context). The used long - context / document - level dataset for Quality Estimation of Machine Translation is an augmented variant of the sentence - level WMT DA Human Evaluation dataset. In addition to individual sentences, it contains augmentations of 2, 4, 8, 16, and 32 sentences, among each language pair lp and domain. The raw column represents a weighted average of scores of augmented sentences using character lengths of src and mt as weights.

• Training data: 7.65 million long - context texts
• Test data: 59,235 long - context texts

Training procedure

The model is trained on 1x H200 SXM (143 GB VRAM) for approx. 26 hours.

• tokenizer.model_max_length: 8192 (full context length)
• attn_implementation: flash_attention_2

Training hyperparameters

The following hyperparameters were used during training:

• learning_rate: 0.0003
• train_batch_size: 128
• eval_batch_size: 128
• seed: 42
• optimizer: Use OptimizerNames.ADAMW_TORCH_FUSED with betas=(0.9,0.999) and epsilon = 1e - 08 and optimizer_args = No additional optimizer arguments
• lr_scheduler_type: linear
• training_steps: 60000 (approx. 1 epoch)

Training results

Framework versions

• Transformers 4.48.1
• Pytorch 2.4.1+cu124
• Datasets 3.2.0
• Tokenizers 0.21.0

🔧 Technical Details

Model Index

Metrics

• pearsonr
• mae
• r_squared

📄 License

This project is licensed under the apache - 2.0 license.