針對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")