T5

In [ ]:

import numpy as np
import torch
from datasets import Dataset,DatasetDict
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, DataCollatorForSeq2Seq, Seq2SeqTrainer, Seq2SeqTrainingArguments
import argparse
from tools.download import save
from rouge import Rouge
from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
from pycocoevalcap.bleu.bleu import Bleu
import pandas as pd
from tools.draw import draw1,draw2
import csv

import numpy as np import torch from datasets import Dataset,DatasetDict from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, DataCollatorForSeq2Seq, Seq2SeqTrainer, Seq2SeqTrainingArguments import argparse from tools.download import save from rouge import Rouge from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction from pycocoevalcap.bleu.bleu import Bleu import pandas as pd from tools.draw import draw1,draw2 import csv

In [ ]:

parser = argparse.ArgumentParser()
parser.add_argument("--lr", type=int, default=1e-4)
parser.add_argument("--epochs", type=int, default=6)
parser.add_argument("--batch_size", type=int, default=32)
parser.add_argument("--device", type=str, default="cuda")
parser.add_argument("--download", type=int, default=0)
args = parser.parse_args()

parser = argparse.ArgumentParser() parser.add_argument("--lr", type=int, default=1e-4) parser.add_argument("--epochs", type=int, default=6) parser.add_argument("--batch_size", type=int, default=32) parser.add_argument("--device", type=str, default="cuda") parser.add_argument("--download", type=int, default=0) args = parser.parse_args()

In [ ]:

device = args.device
if torch.cuda.is_available()==False and device == "cuda":
    device = "cpu"
print(f"{device} is in.")

device = args.device if torch.cuda.is_available()==False and device == "cuda": device = "cpu" print(f"{device} is in.")

In [ ]:

if args.download:
    save()

if args.download: save()

In [ ]:

rouge = Rouge()
results = []
dataset_train = Dataset.from_csv("train.csv")
dataset_test = Dataset.from_csv("test.csv")
dataset_train = dataset_train.train_test_split(test_size=0.2)

rouge = Rouge() results = [] dataset_train = Dataset.from_csv("train.csv") dataset_test = Dataset.from_csv("test.csv") dataset_train = dataset_train.train_test_split(test_size=0.2)

In [ ]:

# 打包到一起，用DatasetDict结构，参考HuggingFace
dataset = DatasetDict({
    "train": dataset_train["train"], 
    "validation": dataset_train["test"], 
    "test": dataset_test
})

# 打包到一起，用DatasetDict结构，参考HuggingFace dataset = DatasetDict({ "train": dataset_train["train"], "validation": dataset_train["test"], "test": dataset_test })

In [ ]:

tokenizer = AutoTokenizer.from_pretrained("./tokenizer_T5")
model = AutoModelForSeq2SeqLM.from_pretrained("./model_T5")

tokenizer = AutoTokenizer.from_pretrained("./tokenizer_T5") model = AutoModelForSeq2SeqLM.from_pretrained("./model_T5")

In [ ]:

def process(x):
    # 使用tokenizer对文本进行处理
    description = ["summarize: \n" + e for e in x["description"]]
    inputs = tokenizer(description, max_length=160, truncation=True)
    labels = tokenizer(text_target=x["diagnosis"], max_length=160, truncation=True)
    inputs["labels"] = labels["input_ids"]
    return inputs

def process(x): # 使用tokenizer对文本进行处理 description = ["summarize: \n" + e for e in x["description"]] inputs = tokenizer(description, max_length=160, truncation=True) labels = tokenizer(text_target=x["diagnosis"], max_length=160, truncation=True) inputs["labels"] = labels["input_ids"] return inputs

In [ ]:

tkdata = dataset.map(process, batched=True)

tkdata = dataset.map(process, batched=True)

In [ ]:

def compute_metrics(pred_labels):
    predictions, labels = pred_labels
    decode_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)
    
    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
    decode_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
    
    decode_preds = [" ".join(p) for p in decode_preds]
    decode_labels = [" ".join(l) for l in decode_labels]
    scores = rouge.get_scores(decode_preds, decode_labels, avg=True)

    rouge_scores = scores["rouge-2"]["r"]
    bleu_scores = np.mean([sentence_bleu([l.split()], p.split(), smoothing_function=SmoothingFunction().method4) for l, p in zip(decode_labels, decode_preds)])
    result = {
        "rouge2r": rouge_scores,
        "bleu4": bleu_scores
    }
    results.append(result)
    return result

def compute_metrics(pred_labels): predictions, labels = pred_labels decode_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True) labels = np.where(labels != -100, labels, tokenizer.pad_token_id) decode_labels = tokenizer.batch_decode(labels, skip_special_tokens=True) decode_preds = [" ".join(p) for p in decode_preds] decode_labels = [" ".join(l) for l in decode_labels] scores = rouge.get_scores(decode_preds, decode_labels, avg=True) rouge_scores = scores["rouge-2"]["r"] bleu_scores = np.mean([sentence_bleu([l.split()], p.split(), smoothing_function=SmoothingFunction().method4) for l, p in zip(decode_labels, decode_preds)]) result = { "rouge2r": rouge_scores, "bleu4": bleu_scores } results.append(result) return result

In [ ]:

training_args = Seq2SeqTrainingArguments(
    output_dir="./T5summary",
    per_device_train_batch_size=args.batch_size,
    per_device_eval_batch_size=4 * args.batch_size,
    gradient_accumulation_steps=4,
    logging_steps=8,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    learning_rate=args.lr,
    num_train_epochs=args.epochs,
    predict_with_generate=True
)

training_args = Seq2SeqTrainingArguments( output_dir="./T5summary", per_device_train_batch_size=args.batch_size, per_device_eval_batch_size=4 * args.batch_size, gradient_accumulation_steps=4, logging_steps=8, evaluation_strategy="epoch", save_strategy="epoch", learning_rate=args.lr, num_train_epochs=args.epochs, predict_with_generate=True )

In [ ]:

trainer = Seq2SeqTrainer(
    args=training_args,
    model=model,
    train_dataset=tkdata["train"],
    eval_dataset=tkdata["validation"],
    compute_metrics=compute_metrics,
    tokenizer=tokenizer,
    data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer)
)
trainer.train()
# 在测试集上看看结果
test_result = trainer.evaluate(tkdata["test"])
print(test_result)

trainer = Seq2SeqTrainer( args=training_args, model=model, train_dataset=tkdata["train"], eval_dataset=tkdata["validation"], compute_metrics=compute_metrics, tokenizer=tokenizer, data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer) ) trainer.train() # 在测试集上看看结果 test_result = trainer.evaluate(tkdata["test"]) print(test_result)

In [ ]:

draw1(trainer.state.log_history)
draw2(results)
fields = set().union(*[entry.keys() for entry in results])
with open("T5results.csv", 'w', newline='') as csvfile:
    csvwriter = csv.DictWriter(csvfile, fieldnames=fields)
    csvwriter.writeheader()
    csvwriter.writerows(results)

draw1(trainer.state.log_history) draw2(results) fields = set().union(*[entry.keys() for entry in results]) with open("T5results.csv", 'w', newline='') as csvfile: csvwriter = csv.DictWriter(csvfile, fieldnames=fields) csvwriter.writeheader() csvwriter.writerows(results)