BART

In [1]:

import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import numpy as np
from transformers import AutoTokenizer,AutoModelForSeq2SeqLM,DataCollatorForSeq2Seq,Seq2SeqTrainer,Seq2SeqTrainingArguments,EncoderDecoderModel
import nltk
from sklearn.model_selection import train_test_split
from rouge import Rouge
from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
import matplotlib.pyplot as plt

import evaluate

import pandas as pd import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader import numpy as np from transformers import AutoTokenizer,AutoModelForSeq2SeqLM,DataCollatorForSeq2Seq,Seq2SeqTrainer,Seq2SeqTrainingArguments,EncoderDecoderModel import nltk from sklearn.model_selection import train_test_split from rouge import Rouge from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction import matplotlib.pyplot as plt import evaluate

In [2]:

class MyDataset(Dataset):
    def __init__(self, description, diagnosis, tokenizer, max_len):
        self.description = description
        self.diagnosis = diagnosis
        self.tokenizer = tokenizer
        self.max_len = max_len

    def __len__(self):
        return len(self.description)

    def __getitem__(self, idx):
        description_encoding = self.tokenizer(
            self.description[idx],
            add_special_tokens=True,
            max_length=self.max_len,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )
        diagnosis_encoding = self.tokenizer(
            self.diagnosis[idx],
            add_special_tokens=True,
            max_length=self.max_len,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )
        # 紧凑化
        return {
            'input_ids': description_encoding['input_ids'].flatten(),
            'attention_mask': description_encoding['attention_mask'].flatten(),
            'labels': diagnosis_encoding['input_ids'].flatten().tolist()
        }

class MyDataset(Dataset): def __init__(self, description, diagnosis, tokenizer, max_len): self.description = description self.diagnosis = diagnosis self.tokenizer = tokenizer self.max_len = max_len def __len__(self): return len(self.description) def __getitem__(self, idx): description_encoding = self.tokenizer( self.description[idx], add_special_tokens=True, max_length=self.max_len, padding='max_length', truncation=True, return_tensors='pt' ) diagnosis_encoding = self.tokenizer( self.diagnosis[idx], add_special_tokens=True, max_length=self.max_len, padding='max_length', truncation=True, return_tensors='pt' ) # 紧凑化 return { 'input_ids': description_encoding['input_ids'].flatten(), 'attention_mask': description_encoding['attention_mask'].flatten(), 'labels': diagnosis_encoding['input_ids'].flatten().tolist() }

In [3]:

train_data = pd.read_csv('train.csv')
test_data = pd.read_csv('test.csv')

train_data = pd.read_csv('train.csv') test_data = pd.read_csv('test.csv')

In [4]:

# 划分训练集和验证集
X = train_data['description'].values
y = train_data['diagnosis'].values
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, shuffle=True, random_state=42)

# 划分训练集和验证集 X = train_data['description'].values y = train_data['diagnosis'].values X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, shuffle=True, random_state=42)

In [5]:

tokenizer = AutoTokenizer.from_pretrained("./tokenizer_bartbase")

tokenizer = AutoTokenizer.from_pretrained("./tokenizer_bartbase")

In [6]:

model = AutoModelForSeq2SeqLM.from_pretrained("./model_bartbase")

model = AutoModelForSeq2SeqLM.from_pretrained("./model_bartbase")

In [7]:

MAX_LEN = max([len(i.split(' ')) for i in X])
BATCH_SIZE = 32
print(MAX_LEN)
MAX_LEN1 = max([len(i.split(' ')) for i in test_data['description'].values])
print(MAX_LEN1)

MAX_LEN = max([len(i.split(' ')) for i in X]) BATCH_SIZE = 32 print(MAX_LEN) MAX_LEN1 = max([len(i.split(' ')) for i in test_data['description'].values]) print(MAX_LEN1)

148
145

In [8]:

train_dataset = MyDataset(X_train, y_train, tokenizer, MAX_LEN)
val_dataset = MyDataset(X_test, y_test, tokenizer, MAX_LEN)
test_dataset = MyDataset(test_data['description'].values, test_data['diagnosis'].values, tokenizer, MAX_LEN)

train_dataset = MyDataset(X_train, y_train, tokenizer, MAX_LEN) val_dataset = MyDataset(X_test, y_test, tokenizer, MAX_LEN) test_dataset = MyDataset(test_data['description'].values, test_data['diagnosis'].values, tokenizer, MAX_LEN)

In [9]:

device = "cuda"
if torch.cuda.is_available():
    device = torch.device("cuda")
else:
    device = torch.device("cpu")
print(device)

device = "cuda" if torch.cuda.is_available(): device = torch.device("cuda") else: device = torch.device("cpu") print(device)

cuda

In [10]:

results = []
rouge1 = Rouge()

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 = rouge1.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

results = [] rouge1 = Rouge() 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 = rouge1.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 [11]:

training_args = Seq2SeqTrainingArguments(
    output_dir="./BARTsummary",
    per_device_train_batch_size=32,
    per_device_eval_batch_size=128,
    gradient_accumulation_steps=4,
    logging_steps=8,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    learning_rate=1e-4,
    num_train_epochs=10,
    predict_with_generate=True
)

training_args = Seq2SeqTrainingArguments( output_dir="./BARTsummary", per_device_train_batch_size=32, per_device_eval_batch_size=128, gradient_accumulation_steps=4, logging_steps=8, evaluation_strategy="epoch", save_strategy="epoch", learning_rate=1e-4, num_train_epochs=10, predict_with_generate=True )

In [12]:

trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset,
    tokenizer=tokenizer,
    compute_metrics=compute_metrics,
    data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer)
)

trainer = Seq2SeqTrainer( model=model, args=training_args, train_dataset=train_dataset, eval_dataset=val_dataset, tokenizer=tokenizer, compute_metrics=compute_metrics, data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer) )

In [13]:

trainer.train()

trainer.train()

You're using a BartTokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.

[1120/1120 46:00, Epoch 9/10]

Epoch	Training Loss	Validation Loss	Rouge2r	Bleu4
0	0.484900	0.443210	0.433229	0.203408
2	0.382500	0.361924	0.507505	0.275782
2	0.353400	0.337537	0.534387	0.304564
4	0.344400	0.324195	0.545088	0.313879
4	0.318700	0.316232	0.547460	0.321791
6	0.311700	0.311865	0.547064	0.322491
6	0.303700	0.307329	0.552348	0.328163
8	0.289400	0.304591	0.553974	0.329657
8	0.275900	0.303300	0.552707	0.330076
9	0.289700	0.302620	0.556572	0.332876

Out[13]:

TrainOutput(global_step=1120, training_loss=0.5145041831902095, metrics={'train_runtime': 2762.5936, 'train_samples_per_second': 52.125, 'train_steps_per_second': 0.405, 'total_flos': 1.26337394147328e+16, 'train_loss': 0.5145041831902095, 'epoch': 9.96})

In [14]:

# 在测试集上看看结果
test_result = trainer.evaluate(test_dataset)
print(test_result)

# 在测试集上看看结果 test_result = trainer.evaluate(test_dataset) print(test_result)

[16/16 00:39]

{'eval_loss': 0.30247724056243896, 'eval_rouge2r': 0.5642511549283902, 'eval_bleu4': 0.3451178296802012, 'eval_runtime': 44.5588, 'eval_samples_per_second': 44.885, 'eval_steps_per_second': 0.359, 'epoch': 9.96}

In [15]:

# 损失函数与epochs变化情况 在训练集和验证集上
history1 = trainer.state.log_history

epochs = [entry.get('epoch', 0) for entry in history1]
train_losses = [entry.get('loss', None) if 'loss' in entry else None for entry in history1]
eval_losses = [entry.get('eval_loss', None) if 'eval_loss' in entry else None for entry in history1]

data = [(epoch, train_loss, eval_loss) for epoch, train_loss, eval_loss in zip(epochs, train_losses, eval_losses) if train_loss is not None or eval_loss is not None]
epochs, train_losses, eval_losses = zip(*data) if data else ([], [], [])

# 创建折线图
plt.plot(epochs, train_losses, marker='o', label='Train', color='blue')
plt.plot(epochs, eval_losses, marker='X', label='Eval', color='red')


plt.title('Loss Change')
plt.xlabel('Epoch')
plt.ylabel('Loss')

plt.grid(True)

plt.legend()
plt.show()

# 损失函数与epochs变化情况 在训练集和验证集上 history1 = trainer.state.log_history epochs = [entry.get('epoch', 0) for entry in history1] train_losses = [entry.get('loss', None) if 'loss' in entry else None for entry in history1] eval_losses = [entry.get('eval_loss', None) if 'eval_loss' in entry else None for entry in history1] data = [(epoch, train_loss, eval_loss) for epoch, train_loss, eval_loss in zip(epochs, train_losses, eval_losses) if train_loss is not None or eval_loss is not None] epochs, train_losses, eval_losses = zip(*data) if data else ([], [], []) # 创建折线图 plt.plot(epochs, train_losses, marker='o', label='Train', color='blue') plt.plot(epochs, eval_losses, marker='X', label='Eval', color='red') plt.title('Loss Change') plt.xlabel('Epoch') plt.ylabel('Loss') plt.grid(True) plt.legend() plt.show()

In [16]:

epochs = range(0,len(results))
rouge2r_scores = [entry["rouge2r"] for entry in results]
bleu4_scores = [entry["bleu4"] for entry in results]

plt.plot(epochs, rouge2r_scores, marker='o', label='ROUGE-2 Recall', color='blue')
plt.plot(epochs, bleu4_scores, marker='o', label='BLEU-4', color='red')

plt.title('BART Metrics Change over Epochs')
plt.xlabel('Epoch')
plt.ylabel('Scores')

plt.grid(True)
plt.legend()
plt.show()

epochs = range(0,len(results)) rouge2r_scores = [entry["rouge2r"] for entry in results] bleu4_scores = [entry["bleu4"] for entry in results] plt.plot(epochs, rouge2r_scores, marker='o', label='ROUGE-2 Recall', color='blue') plt.plot(epochs, bleu4_scores, marker='o', label='BLEU-4', color='red') plt.title('BART Metrics Change over Epochs') plt.xlabel('Epoch') plt.ylabel('Scores') plt.grid(True) plt.legend() plt.show()

In [17]:

import csv
fields = set().union(*[entry.keys() for entry in results])
with open("BARTresults.csv", 'w', newline='') as csvfile:
    csvwriter = csv.DictWriter(csvfile, fieldnames=fields)
    csvwriter.writeheader()
    csvwriter.writerows(results)

import csv fields = set().union(*[entry.keys() for entry in results]) with open("BARTresults.csv", 'w', newline='') as csvfile: csvwriter = csv.DictWriter(csvfile, fieldnames=fields) csvwriter.writeheader() csvwriter.writerows(results)

In [ ]: