Main

In [ ]:

from gru import GRU
import T5
import torch
import pandas as pd
import numpy as np
import torch.optim as optim
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from tools.download import save
from tools.load import load,load_easy
from tools.train import train,valid,rouge
from tools.draw import draw1,draw2
from transformers import AutoTokenizer,AutoModelForSeq2SeqLM,DataCollatorForSeq2Seq,Seq2SeqTrainer,Seq2SeqTrainingArguments,EncoderDecoderModel
import argparse
import nltk
from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
from sklearn.model_selection import train_test_split
from rouge import Rouge
import matplotlib.pyplot as plt
import evaluate
import csv

from gru import GRU import T5 import torch import pandas as pd import numpy as np import torch.optim as optim import torch.nn as nn from torch.utils.data import Dataset, DataLoader from tools.download import save from tools.load import load,load_easy from tools.train import train,valid,rouge from tools.draw import draw1,draw2 from transformers import AutoTokenizer,AutoModelForSeq2SeqLM,DataCollatorForSeq2Seq,Seq2SeqTrainer,Seq2SeqTrainingArguments,EncoderDecoderModel import argparse import nltk from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction from sklearn.model_selection import train_test_split from rouge import Rouge import matplotlib.pyplot as plt import evaluate import csv

In [ ]:

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 [ ]:

parser = argparse.ArgumentParser()
parser.add_argument("--model", type=str, default="BART")
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("--model", type=str, default="BART") 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 [ ]:

if args.model == "BART":
    tokenizer = AutoTokenizer.from_pretrained("./tokenizer_bartbase")
    model = AutoModelForSeq2SeqLM.from_pretrained("./model_bartbase")
    train_dataset, val_dataset, test_dataset = load_easy(tokenizer, 148)
    training_args = Seq2SeqTrainingArguments(
        output_dir="./BARTsummary",
        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
    )
    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.train()
    # 在测试集上看看结果
    test_result = trainer.evaluate(test_dataset)
    print(test_result)
    draw1(trainer.state.log_history)
    draw2(results)
    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)

if args.model == "BART": tokenizer = AutoTokenizer.from_pretrained("./tokenizer_bartbase") model = AutoModelForSeq2SeqLM.from_pretrained("./model_bartbase") train_dataset, val_dataset, test_dataset = load_easy(tokenizer, 148) training_args = Seq2SeqTrainingArguments( output_dir="./BARTsummary", 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 ) 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.train() # 在测试集上看看结果 test_result = trainer.evaluate(test_dataset) print(test_result) draw1(trainer.state.log_history) draw2(results) 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 [ ]:

elif args.model == "GRU":
    train_loader, val_loader, test_loader = load(batch_size=64)
    model=GRU(
        num_hiddens=256, 
        embedding_size=256,
        encoder_layers=2,
        decoder_layers=2, 
        direction=2,
        vocab_size=1000
    )
    loss_func = nn.CrossEntropyLoss(reduction="sum")
    optimizer = optim.AdamW(model.parameters(), lr=args.lr)
    train(
        model=model, 
        loss_func=loss_func, 
        optimizer=optimizer,
        data_loader=train_loader,
        device=device, 
        epochs=args.epochs
    )
    valid(model, test_loader, device)

elif args.model == "GRU": train_loader, val_loader, test_loader = load(batch_size=64) model=GRU( num_hiddens=256, embedding_size=256, encoder_layers=2, decoder_layers=2, direction=2, vocab_size=1000 ) loss_func = nn.CrossEntropyLoss(reduction="sum") optimizer = optim.AdamW(model.parameters(), lr=args.lr) train( model=model, loss_func=loss_func, optimizer=optimizer, data_loader=train_loader, device=device, epochs=args.epochs ) valid(model, test_loader, device)