Modernbert Base Long Context Qe V1

🚀 机器翻译质量评估模型

本模型用于机器翻译的质量评估，无需参考文本，适用于长上下文场景，能有效提升文档级翻译质量评估的准确性。

🚀 快速开始

安装所需库

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

加载测试数据集

from datasets import load_dataset

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

加载模型和分词器

from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

# 加载微调后的模型和分词器
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)

# 若GPU可用，将模型移至GPU
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
model.eval()

准备数据集

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

生成预测

基础用法

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]

高级用法

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

# 分词函数
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
    


# 创建DataLoader进行批处理
test_dataloader = DataLoader(test_dataset, 
                             batch_size=128,   # 根据需要调整批大小
                             shuffle=False)


# 存储所有预测结果的列表
predictions = []

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

        tokens = process_batch(batch, tokenizer, device)
        
        # 前向传播：生成模型的logits
        outputs = model(**tokens)

        # 获取logits（预测结果）
        logits = outputs.logits

        # 提取回归预测值
        batch_predictions = logits.squeeze()

        # 将预测结果添加到列表中
        predictions.extend(batch_predictions.tolist())

✨ 主要特性

• 无参考长上下文评估：可对机器翻译系统进行无参考的长上下文质量评估，适用于文档级质量评估。
• 多语言支持：支持多种语言，包括但不限于英语、中文、日语、俄语等。

📦 安装指南

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

📚 详细文档

模型描述

本模型是在 answerdotai/ModernBERT-base 基础上，针对 ymoslem/wmt-da-human-evaluation-long-context 数据集进行微调得到的。用于机器翻译的无参考长上下文质量评估，训练数据包含最多32个句子的翻译对，适合文档级质量评估。

训练和评估数据

• 训练数据：765万条长上下文文本。
• 测试数据：59,235条长上下文文本。

训练过程

训练超参数

• 学习率：0.0003
• 训练批次大小：128
• 评估批次大小：128
• 随机种子：42
• 优化器：使用 OptimizerNames.ADAMW_TORCH_FUSED，betas=(0.9,0.999)，epsilon=1e-08
• 学习率调度器类型：线性
• 训练步数：60000（约1个epoch）

训练结果

框架版本

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

推理

按照上述“快速开始”部分的步骤进行推理，可使用 text-classification 管道或详细代码进行预测。

🔧 技术细节

• 基础模型：answerdotai/ModernBERT-base
• 训练设备：1x H200 SXM（143 GB VRAM），训练约26小时。
• 分词器最大长度：8192（全上下文长度）
• 注意力实现：flash_attention_2

📄 许可证

本模型使用 Apache-2.0 许可证。

📊 模型指标