File size: 2,388 Bytes
ca4fc4d |
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 |
import matplotlib.pyplot as plt
import time
import torch
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import numpy as np
import tracemalloc
# from Andromeda.model import Andromeda
from Andromeda.model import Andromeda
from Andromeda.utils.stable_adamw import StableAdamWUnfused
torch.manual_seed(0)
if torch.cuda.is_available():
torch.cuda.manual_seed(0)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
import torch.nn.functional as F
from nltk.translate.bleu_score import corpus_bleu
from rouge import Rouge
from sklearn.metrics import f1_score
class AccuracyMetrics:
def __init__(self):
self.rouge = Rouge()
def calculate_perplexity(self, model, data_loader):
model.eval()
total_loss = 0
with torch.no_grad():
for batch in data_loader:
input_ids, labels = batch
output = model(input_ids)
loss = F.cross_entropy(output.view(-1, output.size(-1)), labels.view(-1))
total_loss += loss.item()
return torch.exp(torch.tensor(total_loss / len(data_loader)))
def calculate_bleu(self, references, hypotheses):
return corpus_bleu(references, hypotheses)
def calculate_rouge(self, references, hypotheses):
scores = self.rouge.get_scores(hypotheses, references, avg=True)
return scores
def calculate_f1(self, true_labels, pred_labels):
return f1_score(true_labels, pred_labels, average="weighted")
#mock test dataset
test_dataset = datasets.FakeData(size=1000, transform=transforms.ToTensor())
#model
model = Andromeda(
num_tokens=50304,
dim=1024,
depth=24,
dim_head=128,
heads=8,
alibi_num_heads=4
)
# Usage:
accuracy_metrics = AccuracyMetrics()
# Calculate Perplexity
perplexity = accuracy_metrics.calculate_perplexity(model, data_loader)
print('Perplexity:', perplexity)
# Calculate BLEU
bleu = accuracy_metrics.calculate_bleu(references, hypotheses)
print('BLEU Score:', bleu)
# Calculate ROUGE
rouge_scores = accuracy_metrics.calculate_rouge(references, hypotheses)
print('ROUGE Scores:', rouge_scores)
# Calculate F1 Score
f1 = accuracy_metrics.calculate_f1(true_labels, pred_labels)
print('F1 Score:', f1)
# Add at the bottom of your file
if __name__ == "__main__":
AccuracyMetrics() |