File size: 3,667 Bytes
14e6deb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 |
#Helsinki-NLP/opus-mt-zh-en
# 测试中英翻译模型
# from transformers import pipeline
# translator = pipeline("translation", model="Helsinki-NLP/opus-mt-en-zh", max_time=7)
# prediction = translator("FRST", )[0]["translation_text"]
# print(prediction)
# 微调
from datasets import load_dataset, load_metric
import torch
import numpy as np
import os
raw_datasets = load_dataset("json", data_files="./more models/bank_en_zh_4.json")
split_datasets = raw_datasets["train"].train_test_split(train_size=0.9, seed=20)
split_datasets["validation"] = split_datasets.pop("test")
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, DataCollatorForSeq2Seq
model_checkpoint = "Helsinki-NLP/opus-mt-en-zh"
# tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, return_tensors="tf",device=device)
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, return_tensors="tf")
model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)
# model = torch.nn.DataParallel(model)
model.cuda()
# os.environ['CUDA_VISIBLE_DEVICES'] = '0'
# device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# model.to(device)
max_input_length = 23
max_target_length = 23
def preprocess_function(examples):
inputs = [ex["en"] for ex in examples["translation"]]
targets = [ex["zh"] for ex in examples["translation"]]
model_inputs = tokenizer(inputs, max_length=max_input_length, truncation=True)
# Set up the tokenizer for targets
with tokenizer.as_target_tokenizer():
labels = tokenizer(targets, max_length=max_target_length, truncation=True)
model_inputs["labels"] = labels["input_ids"]
return model_inputs
tokenized_datasets = split_datasets.map(
preprocess_function,
batched=True,
remove_columns=split_datasets["train"].column_names,
)
data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
batch = data_collator([tokenized_datasets["train"][i] for i in range(1, 3)])
def compute_metrics(eval_preds):
preds, labels = eval_preds
# In case the model returns more than the prediction logits
if isinstance(preds, tuple):
preds = preds[0]
decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
# Replace -100s in the labels as we can't decode them
labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
# Some simple post-processing
decoded_preds = [pred.strip() for pred in decoded_preds]
decoded_labels = [[label.strip()] for label in decoded_labels]
# print("Return:- ", metric.compute(predictions=decoded_preds, references=decoded_labels))
# print("decoded_preds:- ", decoded_preds)
# print("decoded_labels:- ", decoded_labels)
# print("Done")
return metric.compute(predictions=decoded_preds, references=decoded_labels)
from transformers import Seq2SeqTrainingArguments
args = Seq2SeqTrainingArguments(
f"marian-finetuned-kde4-en-to-zh",
evaluation_strategy="no",
save_strategy="epoch",
learning_rate=2e-5,
per_device_train_batch_size=128,#32
per_device_eval_batch_size=64,#64
weight_decay=0.01,
save_total_limit=3,
num_train_epochs=60,
predict_with_generate=True,
fp16=True,
push_to_hub=False,
)
from transformers import Seq2SeqTrainer
trainer = Seq2SeqTrainer(
model,
args,
train_dataset=tokenized_datasets["train"],
eval_dataset=tokenized_datasets["validation"],
data_collator=data_collator,
tokenizer=tokenizer,
compute_metrics=compute_metrics,
)
trainer.train()
trainer.save_model("./more models/test-ml-trained_4")
|