Spaces:
Runtime error
Runtime error
File size: 4,753 Bytes
6aee98f |
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 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 |
from typing import List, Tuple
import numpy as np
import pytorch_lightning as pl
import sentencepiece as spm
import torch
import torch.nn as nn
#from sacrebleu.metrics.bleu import BLEU, _get_tokenizer
from torch import optim
from torch.nn.init import xavier_uniform_
from transformers.models.bert.tokenization_bert import BertTokenizer
#import wandb
from dataset import Batched, DecodedBatch
#from models.scheduler import WarmupDecayLR
from transformer import Decoder, Encoder
#from kobe.utils import helpers
class KobeModel(pl.LightningModule):
def __init__(self, args):
super(KobeModel, self).__init__()
self.encoder = Encoder(
vocab_size=args.text_vocab_size + args.cond_vocab_size,
max_seq_len=args.max_seq_len,
d_model=args.d_model,
nhead=args.nhead,
num_layers=args.num_encoder_layers,
dropout=args.dropout,
mode=args.mode,
)
self.decoder = Decoder(
vocab_size=args.text_vocab_size,
max_seq_len=args.max_seq_len,
d_model=args.d_model,
nhead=args.nhead,
num_layers=args.num_decoder_layers,
dropout=args.dropout,
)
self.lr = args.lr
self.d_model = args.d_model
self.loss = nn.CrossEntropyLoss(
reduction="mean", ignore_index=0, label_smoothing=0.1
)
self._reset_parameters()
self.decoding_strategy = args.decoding_strategy
self.vocab = BertTokenizer.from_pretrained(args.text_vocab_path)
#self.bleu = BLEU(tokenize=args.tokenize)
#self.sacre_tokenizer = _get_tokenizer(args.tokenize)()
#self.bert_scorer = BERTScorer(lang=args.tokenize, rescale_with_baseline=True)
def _reset_parameters(self):
for p in self.parameters():
if p.dim() > 1:
xavier_uniform_(p)
def _tokenwise_loss_acc(
self, logits: torch.Tensor, batch: Batched
) -> Tuple[torch.Tensor, float]:
unmask = ~batch.description_token_ids_mask.T[1:]
unmasked_logits = logits[unmask]
unmasked_targets = batch.description_token_ids[1:][unmask]
#acc = helpers.accuracy(unmasked_logits, unmasked_targets)
return self.loss(logits.transpose(1, 2), batch.description_token_ids[1:]), 1
def training_step(self, batch: Batched, batch_idx: int):
encoded = self.encoder.forward(batch)
logits = self.decoder.forward(batch, encoded)
loss, acc = self._tokenwise_loss_acc(logits, batch)
self.lr_schedulers().step()
self.log("train/loss", loss.item())
self.log("train/acc", acc)
return loss
def _shared_eval_step(self, batch: Batched, batch_idx: int) -> DecodedBatch:
encoded = self.encoder.forward(batch)
#logits = self.decoder.forward(batch, encoded)
#loss, acc = self._tokenwise_loss_acc(logits, batch)
preds = self.decoder.predict(
encoded_batch=encoded, decoding_strategy=self.decoding_strategy
)
generated = self.vocab.batch_decode(preds.T.tolist(), skip_special_tokens=True)
#print(generated)
return generated
return DecodedBatch(
loss=loss.item(),
acc=acc,
generated=generated,
descriptions=batch.descriptions,
titles=batch.titles,
)
def validation_step(self, batch, batch_idx):
return self._shared_eval_step(batch, batch_idx)
def test_step(self, batch, batch_idx, dataloader_idx=0):
return self._shared_eval_step(batch, batch_idx)
def _shared_epoch_end(self, outputs: List[DecodedBatch], prefix):
loss = np.mean([o.loss for o in outputs])
acc = np.mean([o.acc for o in outputs])
self.log(f"{prefix}/loss", loss)
self.log(f"{prefix}/acc", acc)
print(outputs)
generated = [g for o in outputs for g in o.generated]
references = [r for o in outputs for r in o.descriptions]
titles = [r for o in outputs for r in o.titles]
# Examples
columns = ["Generated", "Reference"]
data = list(zip(generated[:256:16], references[:256:16]))
table = wandb.Table(data=data, columns=columns)
self.logger.experiment.log({f"examples/{prefix}": table})
def validation_epoch_end(self, outputs):
self._shared_epoch_end(outputs, "val")
def test_epoch_end(self, outputs):
self._shared_epoch_end(outputs, "test")
def configure_optimizers(self):
optimizer = optim.AdamW(self.parameters(), lr=self.lr, betas=(0.9, 0.98))
#scheduler = WarmupDecayLR(optimizer, warmup_steps=10000, d_model=self.d_model)
return [optimizer]
|