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

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

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

from transformers import AutoModel , AutoTokenizer,MarianMTModeltext ="你好，你是谁？"
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):# 添加目标语言 tokeninputs = [">>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):# 添加目标语言 tokeninputs = [">>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.predictionslabel_ids = pred.label_idspred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)label_ids[label_ids == -100] = tokenizer.pad_token_idlabel_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):# 添加目标语言 tokeninputs = [">>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):# 添加目标语言 tokeninputs = [">>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.predictionslabel_ids = pred.label_idspred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)label_ids[label_ids == -100] = tokenizer.pad_token_idlabel_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")