提拔翻译质量的机密武器：硬提示与软提示微调剖析(二) - ToB企服应用市场:ToB评测及商务社交产业平台

import torch
import torch.nn as nn
from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
from torch.utils.data import Dataset, DataLoader
# 定义软提示词嵌入
class SoftPrompt(nn.Module):
def __init__(self, hidden_size, prompt_length):
super().__init__()
self.prompt = nn.Parameter(torch.randn(prompt_length, hidden_size))
# 加载预训练模型和tokenizer
model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50")
tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50")
# 初始化软提示词嵌入
prompt_length = 10 # 可以根据需要调整
hidden_size = model.config.d_model
soft_prompt = SoftPrompt(hidden_size, prompt_length)
# 训练过程
optimizer = torch.optim.AdamW(list(model.parameters()) + [soft_prompt.prompt], lr=5e-5)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
soft_prompt.to(device)
# --- Few-shot 学习 ---
class FewShotDataset(Dataset):
def __init__(self, examples, tokenizer, max_length=512):
self.examples = examples
self.tokenizer = tokenizer
self.max_length = max_length
def __len__(self):
return len(self.examples)
def __getitem__(self, idx):
example = self.examples[idx]
source_text = example['source']
target_text = example['target']
# 将输入和输出文本进行tokenization
inputs = self.tokenizer(source_text, return_tensors="pt", padding=True, truncation=True, max_length=self.max_length)
targets = self.tokenizer(target_text, return_tensors="pt", padding=True, truncation=True, max_length=self.max_length)
# 返回输入输出ID
return {'input_ids': inputs['input_ids'].squeeze(), 'labels': targets['input_ids'].squeeze()}
# 定义 few-shot 示例
few_shot_examples = [
{"source": "天气今天非常晴朗。", "target": "The weather is sunny today."},
{"source": "他喜欢打篮球。", "target": "He likes to play basketball."},
{"source": "我们将讨论气候变化问题。", "target": "We will discuss climate change."}
]
# 创建数据集和DataLoader
dataset = FewShotDataset(few_shot_examples, tokenizer)
dataloader = DataLoader(dataset, batch_size=2)
# 训练模型
model.train()
for epoch in range(3): # 3个epoch
for batch in dataloader:
input_ids = batch['input_ids'].to(device)
labels = batch['labels'].to(device)
# 将软提示词嵌入添加到模型输入
input_ids = torch.cat([soft_prompt.prompt.repeat(input_ids.size(0), 1, 1), input_ids], dim=1)
# 进行前向传播和计算损失
outputs = model(input_ids=input_ids, labels=labels)
loss = outputs.loss
# 反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f"Epoch {epoch+1}, Loss: {loss.item()}")
# --- 上下文提示 ---
class ContextPromptDataset(Dataset):
def __init__(self, examples, tokenizer, max_length=512):
self.examples = examples
self.tokenizer = tokenizer
self.max_length = max_length
def __len__(self):
return len(self.examples)
def __getitem__(self, idx):
example = self.examples[idx]
context_text = example['context']
target_text = example['target']
# 将上下文和目标翻译文本进行tokenization
inputs = self.tokenizer(context_text, return_tensors="pt", padding=True, truncation=True, max_length=self.max_length)
targets = self.tokenizer(target_text, return_tensors="pt", padding=True, truncation=True, max_length=self.max_length)
# 返回输入输出ID
return {'input_ids': inputs['input_ids'].squeeze(), 'labels': targets['input_ids'].squeeze()}
# 定义上下文示例
context_examples = [
{"context": "这篇文章讨论了气候变化的影响，包括经济、社会和环境层面的变化。", "target": "This article discusses the impacts of climate change, including economic, social, and environmental changes."},
{"context": "随着科技的发展，人工智能正在逐步改变各个行业的面貌。", "target": "With the development of technology, artificial intelligence is gradually changing the landscape of various industries."}
]
# 创建数据集和DataLoader
dataset = ContextPromptDataset(context_examples, tokenizer)
dataloader = DataLoader(dataset, batch_size=2)
# 训练模型
model.train()
for epoch in range(3): # 3个epoch
for batch in dataloader:
input_ids = batch['input_ids'].to(device)
labels = batch['labels'].to(device)
# 将软提示词嵌入添加到模型输入
input_ids = torch.cat([soft_prompt.prompt.repeat(input_ids.size(0), 1, 1), input_ids], dim=1)
# 进行前向传播和计算损失
outputs = model(input_ids=input_ids, labels=labels)
loss = outputs.loss
# 反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f"Epoch {epoch+1}, Loss: {loss.item()}")
# --- 链式推理（CoT） ---
class CotPromptDataset(Dataset):
def __init__(self, examples, tokenizer, max_length=512):
self.examples = examples
self.tokenizer = tokenizer
self.max_length = max_length
def __len__(self):
return len(self.examples)
def __getitem__(self, idx):
example = self.examples[idx]
cot_prompt = example['cot_prompt']
target_text = example['target']
# 将链式推理提示和目标翻译文本进行tokenization
inputs = self.tokenizer(cot_prompt, return_tensors="pt", padding=True, truncation=True, max_length=self.max_length)
targets = self.tokenizer(target_text, return_tensors="pt", padding=True, truncation=True, max_length=self.max_length)
# 返回输入输出ID
return {'input_ids': inputs['input_ids'].squeeze(), 'labels': targets['input_ids'].squeeze()}
# 定义链式推理示例
cot_examples = [
{"cot_prompt": "请分步推理并翻译以下句子：\n1. 理解句子的含义\n2. 分析句子的结构\n3. 根据目标语言语法进行调整\n4. 生成最终翻译\n待翻译句子：气候变化对全球生态系统产生了深远的影响。",
"target": "Please reason step by step and translate the following sentence:\n1. Understand the meaning of the sentence\n2. Analyze the structure of the sentence\n3. Adjust according to the grammar of the target language\n4. Generate the final translation\nThe sentence to be translated: Climate change has had a profound impact on global ecosystems."}
]
# 创建数据集和DataLoader
dataset = CotPromptDataset(cot_examples, tokenizer)
dataloader = DataLoader(dataset, batch_size=2)
# 训练模型
model.train()
for epoch in range(3): # 3个epoch
for batch in dataloader:
input_ids = batch['input_ids'].to(device)
labels = batch['labels'].to(device)
# 将软提示词嵌入添加到模型输入
input_ids = torch.cat([soft_prompt.prompt.repeat(input_ids.size(0), 1, 1), input_ids], dim=1)
# 进行前向传播和计算损失
outputs = model(input_ids=input_ids, labels=labels)
loss = outputs.loss
# 反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f"Epoch {epoch+1}, Loss: {loss.item()}")
# 保存微调后的模型
model.save_pretrained("cot_prompt_model")
tokenizer.save_pretrained("cot_prompt_model")

复制代码