针对Helsinki-NLP/opus-mt-zh-en模型进行双向互翻的微调

引言
 题目听起来有点怪怪的，但是实际上就是对Helsinki-NLP/opus-mt-en-es模型进行微调。但是这个模型是单向的，只支持中到英的翻译，反之则不行。这样的话，如果要做中英双向互翻就需要两个模型，那模型体积直接大了两倍。尤其是部署到手机上，模型的体积是一个非常重要的考虑因素。于是自己就对这个模型的微调过程进行了一些改动，实现了单个模型进行双向互翻的能力。

原生模型
 这里给出原生模型的使用方法：

from transformers import AutoModel , AutoTokenizer,MarianMTModel

text ="你好，你是谁？"
name ='Helsinki-NLP/opus-mt-zh-en'
tokenizer = AutoTokenizer.from_pretrained(name)
model = MarianMTModel.from_pretrained(name)
input_ids = tokenizer.encode(text, return_tensors="pt")
outputs = model.generate(input_ids)
decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(decoded)

需要改动的地方
 因为涉及到互翻，所以首先要告诉模型翻译的方向，具体就是在文本数据之前加一个目标语言的标识符，比如中翻英，原文“你好，你是谁？”，处理后就是“>>eng<< 你好，你是谁？”，英翻中则是“>>zho<< Hello，who are you?”

 因此就引出了一个问题，词表vocab.json中并没有“>>eng<<”和“>>zho<<”，那么分词就会出现问题。我尝试过两种方法来解决：

• 首先是常规的解决办法，我最开始直接将这两个标识当做新的token加入词表中，最后也能跑通。这里只描述思想，具体的实现在下面的代码中会体现。
• 然后就是我自己想的非常规方法，为啥自己又想了非常规的方法呢，那是因为我在训练好模型之后，要将模型转换为CT2的格式，但是这个转换过程中因为添加了2个新token导致了报错，搞了一圈也没有解决，于是直接把词表中两个极其罕见的token给删除了，用两个语言标识替代，这样既不会对翻译产生大的影响，又能完成模型格式的转换。当然，这是需要先改词表后进行微调，顺序不能反了。

解决办法一
 通过下面的代码微调之后，就能得到一个双向的翻译能力的模型了，使用的方法和原生模型一样，直接加载就能推理了。

import torch
import evaluate
import zhconv
from datasets import load_dataset, Dataset
import sacrebleu
import os
from transformers import (
    AutoTokenizer, MarianMTModel,
    Seq2SeqTrainer, Seq2SeqTrainingArguments,
    DataCollatorForSeq2Seq
)

# 加载 tokenizer，并添加语言标签
tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-zh-en")
special_tokens = [">>eng<<", ">>zho<<"]
tokenizer.add_special_tokens({"additional_special_tokens": special_tokens})

# 加载模型，并扩展嵌入层大小
model = MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-zh-en")
model.resize_token_embeddings(len(tokenizer))

# 设置 token ID
model.config.eos_token_id = tokenizer.eos_token_id
model.config.pad_token_id = tokenizer.pad_token_id

'''
加载 Tatoeba 数据集（中英句对）
这里我使用的是公开的数据集，可以通过下面的代码来加载到本地。加载到本地之后就可以把data_files换成你自己的地址
import kagglehub
alvations_tatoeba_path = kagglehub.dataset_download('alvations/tatoeba')
'''
tatoeba_dataset = load_dataset(
    "csv",
    data_files="./data/tatoeba-sentpairs.tsv",
    delimiter="\t",
    encoding="utf-8",
    split="train"
)

# 过滤中英句对（zh→en 和 en→zh）
zh2en_dataset = tatoeba_dataset.filter(lambda x: x['SRC LANG'] == "cmn" and x['TRG LANG'] == 'eng')
en2zh_dataset = tatoeba_dataset.filter(lambda x: x['SRC LANG'] == "eng" and x['TRG LANG'] == 'cmn')

# 预处理函数：添加语言标签 + 分词
def preprocess_zh2en(batch):
    # 添加目标语言 token
    inputs = [">>eng<< " + x for x in batch['SRC']]
    
    # 可选：进行繁转简
    inputs = [zhconv.convert(x, 'zh-cn') for x in inputs]
    targets = batch['TRG']

    # 编码
    inputs_encoded = tokenizer(inputs, max_length=128, truncation=True, padding="max_length")
    outputs_encoded = tokenizer(targets, max_length=128, truncation=True, padding="max_length" )

    return {
        "input_ids": inputs_encoded["input_ids"],
        "attention_mask": inputs_encoded["attention_mask"],
        "decoder_input_ids": outputs_encoded["input_ids"],
        "decoder_attention_mask": outputs_encoded["attention_mask"],
        "labels": outputs_encoded["input_ids"].copy(),  # labels 通常跟 decoder_input_ids 相同（训练时用于 loss）
    }

def preprocess_en2zh(batch):
    # 添加目标语言 token
    inputs = [">>zho<< " + x for x in batch['SRC']]
    
    # 可选：进行繁转简
    targets = batch['TRG']
    targets = [zhconv.convert(x, 'zh-cn') for x in targets]
    
    # 编码
    inputs_encoded = tokenizer(inputs, max_length=128, truncation=True, padding="max_length")
    outputs_encoded = tokenizer(targets, max_length=128, truncation=True, padding="max_length" )

    return {
        "input_ids": inputs_encoded["input_ids"],
        "attention_mask": inputs_encoded["attention_mask"],
        "decoder_input_ids": outputs_encoded["input_ids"],
        "decoder_attention_mask": outputs_encoded["attention_mask"],
        "labels": outputs_encoded["input_ids"].copy(),  # labels 通常跟 decoder_input_ids 相同（训练时用于 loss）
    }

# 数据清洗 + 映射分词
zh2en_dataset = zh2en_dataset.map(preprocess_zh2en, batched=True)
en2zh_dataset = en2zh_dataset.map(preprocess_en2zh, batched=True)

# 合并中→英和英→中双向数据
combined_dataset = Dataset.from_dict({
    key: zh2en_dataset[key] + en2zh_dataset[key] for key in zh2en_dataset.features
})

# 拆分训练集和测试集
split_dataset = combined_dataset.train_test_split(test_size=0.1, seed=42)
train_dataset = split_dataset["train"]
eval_dataset = split_dataset["test"]

def compute_metrics(pred):
    pred_ids = pred.predictions
    label_ids = pred.label_ids

    pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
    label_ids[label_ids == -100] = tokenizer.pad_token_id
    label_str = tokenizer.batch_decode(label_ids, skip_special_tokens=True)

    bleu = sacrebleu.corpus_bleu(pred_str, [label_str])

    # 保存验证的结果到本地文件，这样可以实时查看微调的效果
    save_dir = "./eval_logs"
    os.makedirs(save_dir, exist_ok=True)
    eval_id = f"step_{trainer.state.global_step}" if hasattr(trainer, "state") else "eval"
    output_file = os.path.join(save_dir, f"pred_vs_ref_{eval_id}.txt")

    with open(output_file, "w", encoding="utf-8") as f:
        for i, (pred, ref) in enumerate(zip(pred_str, label_str)):
            f.write(f"Sample {i + 1}:\n")
            f.write(f"Prediction: {pred}\n")
            f.write(f"Reference : {ref}\n")
            f.write("=" * 50 + "\n")

    return {
        "bleu": bleu.score
    }

# 训练参数
training_args = Seq2SeqTrainingArguments(
    output_dir='./model/marian-zh-en-bidirectional',
    num_train_epochs=30,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    logging_steps=50,
    save_steps=100,
    eval_steps=100,
    eval_strategy="steps",
    predict_with_generate=True,
    save_total_limit=10,
    report_to="tensorboard",  # 启用 TensorBoard 日志记录
    logging_dir='./logs',  # 指定 TensorBoard 日志的保存路径
)

# 构建 Trainer
trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset.with_format("torch"),
    eval_dataset=eval_dataset.with_format("torch"),
    tokenizer=tokenizer,
    data_collator=DataCollatorForSeq2Seq(tokenizer, model=model),
    compute_metrics=compute_metrics
)

# 开始训练
trainer.train(resume_from_checkpoint=False)

# 保存模型和 tokenizer
model.save_pretrained("./model/marian-zh-en-bidirectional")
tokenizer.save_pretrained("./model/marian-zh-en-bidirectional")

解决办法二
 上面是针对大众场景，具体的场景需要做具体的改动。本方法就是根据我的业务场景来修改的。

 方法一训练得到的模型是使用tokenizer来编解码，因为目标语言标识符已经加入到词表里了，所以编解码没问题。但是我转为CT2格式之后，分词使用的是sentencepiece模型，具体就是用source.spm、target.spm分别对中文和英文进行分词，然后根据共享词表转换为token的id。 共享词表中是有语言标识符的，但是sentencepiece模型里却没有添加两个新token，所以就无法识别，导致分词错误。我的做法就是推理的时候先不加目标语言的标识符，先分词，然后手动加上去。这样分词就不会出问题了，然后进行编码就能根据共享词表来编码了。

 还有一个问题就是，输入是中英混合的文本，这样sentencepiece分词器也无法正确识别，一个办法就是将中英文分开，分别进行分词，然后将分词的结果按顺序进行拼接。

 最后，以上都是基于不重新训练分词模型的做法，如果可以重新训练分词模型，那么就不需要搞上面哪些操作了。

import torch
import evaluate
import zhconv
from datasets import load_dataset, Dataset
import sacrebleu
import os
from transformers import (
    AutoTokenizer, MarianMTModel,
    Seq2SeqTrainer, Seq2SeqTrainingArguments,
    DataCollatorForSeq2Seq
)

# 加载 tokenizer
tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-zh-en")

# 加载模型
model = MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-zh-en")

# 设置 token ID
model.config.eos_token_id = tokenizer.eos_token_id
model.config.pad_token_id = tokenizer.pad_token_id

'''
加载 Tatoeba 数据集（中英句对）
这里我使用的是公开的数据集，可以通过下面的代码来加载到本地。加载到本地之后就可以把data_files换成你自己的地址
import kagglehub
alvations_tatoeba_path = kagglehub.dataset_download('alvations/tatoeba')
'''
tatoeba_dataset = load_dataset(
    "csv",
    data_files="./data/tatoeba-sentpairs.tsv",
    delimiter="\t",
    encoding="utf-8",
    split="train"
)

# 过滤中英句对（zh→en 和 en→zh）
zh2en_dataset = tatoeba_dataset.filter(lambda x: x['SRC LANG'] == "cmn" and x['TRG LANG'] == 'eng')
en2zh_dataset = tatoeba_dataset.filter(lambda x: x['SRC LANG'] == "eng" and x['TRG LANG'] == 'cmn')

# 预处理函数：添加语言标签 + 分词
def preprocess_zh2en(batch):
    # 添加目标语言 token
    inputs = [">>eng<< " + x for x in batch['SRC']]
    
    # 可选：进行繁转简
    inputs = [zhconv.convert(x, 'zh-cn') for x in inputs]
    targets = batch['TRG']

    # 编码
    inputs_encoded = tokenizer(inputs, max_length=128, truncation=True, padding="max_length")
    outputs_encoded = tokenizer(targets, max_length=128, truncation=True, padding="max_length" )

    return {
        "input_ids": inputs_encoded["input_ids"],
        "attention_mask": inputs_encoded["attention_mask"],
        "decoder_input_ids": outputs_encoded["input_ids"],
        "decoder_attention_mask": outputs_encoded["attention_mask"],
        "labels": outputs_encoded["input_ids"].copy(),  # labels 通常跟 decoder_input_ids 相同（训练时用于 loss）
    }

def preprocess_en2zh(batch):
    # 添加目标语言 token
    inputs = [">>zho<< " + x for x in batch['SRC']]
    
    # 可选：进行繁转简
    targets = batch['TRG']
    targets = [zhconv.convert(x, 'zh-cn') for x in targets]
    
    # 编码
    inputs_encoded = tokenizer(inputs, max_length=128, truncation=True, padding="max_length")
    outputs_encoded = tokenizer(targets, max_length=128, truncation=True, padding="max_length" )

    return {
        "input_ids": inputs_encoded["input_ids"],
        "attention_mask": inputs_encoded["attention_mask"],
        "decoder_input_ids": outputs_encoded["input_ids"],
        "decoder_attention_mask": outputs_encoded["attention_mask"],
        "labels": outputs_encoded["input_ids"].copy(),  # labels 通常跟 decoder_input_ids 相同（训练时用于 loss）
    }

# 数据清洗 + 映射分词
zh2en_dataset = zh2en_dataset.map(preprocess_zh2en, batched=True)
en2zh_dataset = en2zh_dataset.map(preprocess_en2zh, batched=True)

# 合并中→英和英→中双向数据
combined_dataset = Dataset.from_dict({
    key: zh2en_dataset[key] + en2zh_dataset[key] for key in zh2en_dataset.features
})

# 拆分训练集和测试集
split_dataset = combined_dataset.train_test_split(test_size=0.1, seed=42)
train_dataset = split_dataset["train"]
eval_dataset = split_dataset["test"]

def compute_metrics(pred):
    pred_ids = pred.predictions
    label_ids = pred.label_ids

    pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
    label_ids[label_ids == -100] = tokenizer.pad_token_id
    label_str = tokenizer.batch_decode(label_ids, skip_special_tokens=True)

    bleu = sacrebleu.corpus_bleu(pred_str, [label_str])

    # 保存验证的结果到本地文件，这样可以实时查看微调的效果
    save_dir = "./eval_logs"
    os.makedirs(save_dir, exist_ok=True)
    eval_id = f"step_{trainer.state.global_step}" if hasattr(trainer, "state") else "eval"
    output_file = os.path.join(save_dir, f"pred_vs_ref_{eval_id}.txt")

    with open(output_file, "w", encoding="utf-8") as f:
        for i, (pred, ref) in enumerate(zip(pred_str, label_str)):
            f.write(f"Sample {i + 1}:\n")
            f.write(f"Prediction: {pred}\n")
            f.write(f"Reference : {ref}\n")
            f.write("=" * 50 + "\n")

    return {
        "bleu": bleu.score
    }

# 训练参数
training_args = Seq2SeqTrainingArguments(
    output_dir='./model/marian-zh-en-bidirectional',
    num_train_epochs=30,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    logging_steps=50,
    save_steps=100,
    eval_steps=100,
    eval_strategy="steps",
    predict_with_generate=True,
    save_total_limit=10,
    report_to="tensorboard",  # 启用 TensorBoard 日志记录
    logging_dir='./logs',  # 指定 TensorBoard 日志的保存路径
)

# 构建 Trainer
trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset.with_format("torch"),
    eval_dataset=eval_dataset.with_format("torch"),
    tokenizer=tokenizer,
    data_collator=DataCollatorForSeq2Seq(tokenizer, model=model),
    compute_metrics=compute_metrics
)

# 开始训练
trainer.train(resume_from_checkpoint=False)

# 保存模型和 tokenizer
model.save_pretrained("./model/marian-zh-en-bidirectional")
tokenizer.save_pretrained("./model/marian-zh-en-bidirectional")

基于训练好的模型我还搞了一套使用C++来推理的代码，方面在更多的平台使用，具体可以在github上搜"xinliu9451/Opus-Mt_Bidirectional_Translation"。