csc449 · January 4, 2025 23:48
diff --git a/gistfile1.txt b/gistfile1.txt
 #%% imports
 from dataclasses import dataclass
 import math 
 import numpy as np
 import os
 import requests
 import pandas as pd
 import unicodedata
 import re
 import torch
 import torch.nn as nn
 from torch.nn import functional as F
 from torch.utils.data import Dataset, DataLoader

 #%%  download data
 def download_file(url, filename):  
  try:
    response = requests.get(url)
    response.raise_for_status()  # Raise an exception for bad status codes

    with open(filename, 'wb') as f:
      f.write(response.content)


    print(f"File downloaded successfully to '{filename}'")

  except requests.exceptions.RequestException as e:
    print(f"Error downloading file: {e}")

 def unicode_to_ascii(s):
    return ''.join(c for c in unicodedata.normalize('NFD', s) if unicodedata.category(c) != 'Mn')

 def preprocess_sentence(w):
  """
  string normalization, space handling, punctuation handling, unicode, add special tokens.
  """
  w = unicode_to_ascii(w.lower().strip())

  # creating a space between a word and the punctuation following it
  # eg: "he is a boy." => "he is a boy ."
  w = re.sub(r"([?.!,¿])", r" \1 ", w)  # space before punctuation
  w = re.sub(r'[" "]+', " ", w) # multiple space -> one space
  w = re.sub(r"[^a-zA-Z?.!,¿]+", " ", w) # remove unrecognized char, replacing everything with space except (a-z, A-Z, ".", "?", "!", ",")
  w = w.rstrip().strip()

  # defer to later.
  # adding a start and an end token to the sentence, model know when to start and stop predicting.
  # w = '<bos> ' + w + ' <eos>'
  return w
 #endregion

 data_url = 'https://raw.githubusercontent.com/chaoyu-io/yshare/refs/heads/main/data/spa.txt'
 file_path = '/tmp/spa.txt'

 if not os.path.exists(file_path):
    download_file(data_url, file_path) 

 #%%
 with open(file_path, encoding='UTF-8') as fid:
  lines = fid.read().strip().split('\n')

 # sample size (try with smaller sample size to reduce computation)
 PAD_IDX = 50256 # reuse EOS, to avoid embedding look up error
 BOS_IDX = 1
 EOS_IDX = 2
 BLK_SIZE = 1024
 num_examples = 30000
 # creates lists containing each pair
 original_word_pairs = [[w for w in l.split('\t')][:2] for l in lines[:num_examples]] # [:2], first two elements as a pair, as in original data format.

 data = pd.DataFrame(original_word_pairs, columns=["en", "es"])

 # Converts the unicode file to ascii
 data["en"] = data.en.apply(lambda w: preprocess_sentence(w))
 data["es"] = data.es.apply(lambda w: preprocess_sentence(w))

 # %% build vocabulary 
 class LanguageIndex():
    def __init__(self, lang):
       self.lang = lang 
       self.word2idx = {}
       self.idx2word = {}
       self.vocab = set()        
       self.create_index()
    
    def create_index(self):
        for phrase in self.lang:
            self.vocab.update(phrase.split(' '))
        self.vocab = sorted(self.vocab)

        # use -1 ad padding index, in 
        self.word2idx['<pad>'] = PAD_IDX
        self.word2idx['<bos>'] = BOS_IDX
        self.word2idx['<eos>'] = EOS_IDX
        for index,word in enumerate(self.vocab):
            self.word2idx[word] = index + 3 # 3 special tokens.
        
        for word,index in self.word2idx.items():
            self.idx2word[index] = word

 # Vocabulary, convert to index.
 src_vocab = LanguageIndex( data.en.values.tolist())
 tgt_vocab = LanguageIndex( data.es.values.tolist())
 src_tensor_0 = [ [ src_vocab.word2idx[w] for w in s.split(' ')] for s in data.en.values.tolist() ]
 tgt_tensor_0 =  [ [ tgt_vocab.word2idx[w] for w in s.split(' ')] for s in data.es.values.tolist() ]

 import tiktoken
 tokenizer = tiktoken.get_encoding("gpt2")
 EOS_IDX = tokenizer.encode("<|endoftext|>", allowed_special={"<|endoftext|>"})[0]
 print('eos from tiktoken ', EOS_IDX)
 src_tensor = [tokenizer.encode(s) for s in data.en.values.tolist()]
 tgt_tensor = [tokenizer.encode(s) for s in data.es.values.tolist()]

 # max_src_len = max_tensor_len(src_tensor)
 # max_tgt_len = max_tensor_len(tgt_tensor)
 # src_padded = [ pad_tensor(t, max_src_len) for t in src_tensor ]
 # tgt_padded = [pad_tensor(t, max_tgt_len) for t in tgt_tensor]

 # for i,row in data.iterrows():

 # 1) __getitem__ 2) __len__
 class TranslationData(Dataset):
    def __init__(self, src, tgt):
        self.src = src
        self.tgt = tgt
        # self.src_valid_len = [ np.sum( x!=0 ) for x in src]

    def __getitem__(self, index):
        x = self.src[index]
        y = self.tgt[index]
        x_len = len(x)
        y_len = len(y)
        sentence = [EOS_IDX] + x + [EOS_IDX] + y + [EOS_IDX]
        target = sentence[1:] + [PAD_IDX] # extra pad        
        for i in range(x_len+2):
          target[i] = PAD_IDX # mask out (x_len+2), the first prediction align wit y
        return sentence, target, x_len, y_len

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

 # create all data and split
 from torch.utils.data import random_split

 all_data = TranslationData(src_tensor, tgt_tensor)
 assert len(all_data) == num_examples, 'data size match'
 num_train_sample = int(0.8 * num_examples)
 num_test_sample = num_examples - num_train_sample
 train_dataset, test_dataset = random_split(all_data, [num_train_sample, num_test_sample])

 def pad_array(t, max_len):
  # min(max_len): batch data, and model context  window BLK_SIZE
  padded = np.full(( min(max_len, BLK_SIZE)), PAD_IDX, dtype=np.int64) # np.full with init.
  l = min(len(t), max_len) # max_len should always > 5, sanity check.
  padded[:l] = t
  return padded

 def make_batch(data_batch):
  max_l = max( len(src) for src,_,_,_ in data_batch )
  
  all_src_pad = []
  all_tgt_pad = []
  for (src, tgt, x_len, y_len) in data_batch:
    # first target tensor is BOS
    src_pad = pad_array(src, max_l)
    tgt_pad = pad_array(tgt, max_l)
    all_src_pad.append(torch.tensor(src_pad))
    all_tgt_pad.append(torch.tensor(tgt_pad))
  src_tensor = torch.stack(all_src_pad)
  tgt_tensor = torch.stack(all_tgt_pad)
  return (src_tensor, tgt_tensor)

 train_dataloader = DataLoader(train_dataset, batch_size=64, collate_fn=make_batch, shuffle=True)
 test_dataloader = DataLoader(test_dataset, batch_size=64,  collate_fn=make_batch, shuffle=True)

 # %% build transformer

 class LayerNorm(nn.Module):
    def __init__(self, ndim, bias):
        super().__init__()
        # not linear layer, simply nn.Param, initialization built-in
        self.weight = nn.Parameter(torch.ones(ndim)) # init. as 1
        self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None

    def forward(self, input):
        # https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html#torch.nn.LayerNorm
        # layer_norm applied on last D dimensiosn, D is shape of [self.weight.shape]
        # weight, and bias are additionally applied as multiplication, addition after normalization., epsilon is for normalization.
        return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)

 class CausalSelfAttention(nn.Module):

    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        # key, query, value projections for all heads, but in a batch
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias) #(in_features, out_features)
        # output projection
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
        # regularization
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)
        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.dropout = config.dropout
        # flash attention make GPU go brrrrr but support is only in PyTorch >= 2.0
        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
        if not self.flash:
            print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
            # causal mask to ensure that attention is only applied to the left in the input sequence
            # register_buffer: not learneable, save model-specific state information.
            self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size)) # block_size: context length.
                                        .view(1, 1, config.block_size, config.block_size)) # view: reshape.                                        

    def forward(self, x):
        B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)

        # nn.Linear applies to last dimension of inputs: [B,T,C]  nn.Linear(C, 3C), output [B,T,3C], after the split, we got three tensors of [B,T,C]
        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
        q, k, v  = self.c_attn(x).split(self.n_embd, dim=2) # split into 3 sub-matrixese, c_attn [n_embd, 3xn_mebed]
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs) # C got divided.
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)

        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
        if self.flash:
            # efficient attention using Flash Attention CUDA kernels
            y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 0, is_causal=True)
        else:
            # manual implementation of attention
            # @: matmul, q.matmul(k), 
            # matmul applies to last two dimension, [B, nh, T, hs] @ [B, nh, hs, T], --> [B, nh, T, T]
            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) # size(-1) is hs
            # softmax fill with '-inf', masked_fill, if bias==0, set as -inf
            att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf')) # up to :T, T not guaranted to be config.block_size             
            att = F.softmax(att, dim=-1)
            att = self.attn_dropout(att) # apply dropout in the end.
            y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
        # tranpose(1,2) -> [B, T, nh, hs], contiguous ensures continuous memory, following order of dimension, 
        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side

        # output projection
        y = self.resid_dropout(self.c_proj(y))
        return y

 class MLP(nn.Module):

    def __init__(self, config):
        super().__init__()
        # MLP 4-8 times of embedding dimension, larger than transformer hidden.
        self.c_fc    = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
        self.gelu    = nn.GELU()
        self.c_proj  = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        # fully connected, activation, fully connected, dropuot.
        x = self.c_fc(x)
        x = self.gelu(x)
        x = self.c_proj(x)
        x = self.dropout(x)
        return x

 class Block(nn.Module):

    def __init__(self, config):
        super().__init__()
        self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)
        self.mlp = MLP(config)

    def forward(self, x):
        # layer-norm, attention, 
        # res-add 
        x = x + self.attn(self.ln_1(x))
        # layer-norm, mlp 
        # res-add
        x = x + self.mlp(self.ln_2(x))
        return x

 @dataclass
 class GPTConfig:
    # 1024 / 50304
    # 12 / 12 / 768 as GPT-2
    block_size: int = BLK_SIZE  # 1024
    vocab_size: int = 50304 # GPT-2 vocab_size of 50257, padded up to nearest multiple of 64 for efficiency
    n_layer: int = 12 # 12
    n_head: int = 12 # 12
    n_embd: int = 768 # 768
    dropout: float = 0.0
    bias: bool = True # True: bias in Linears and LayerNorms, like GPT-2. False: a bit better and faster

 class GPT(nn.Module):

    def __init__(self, config):
        super().__init__()
        assert config.vocab_size is not None
        assert config.block_size is not None
        self.config = config

        self.transformer = nn.ModuleDict(dict(
            # wte: saves the embedding table, each row corresponding to embedding for token.
            wte = nn.Embedding(config.vocab_size, config.n_embd), # wte: word token embedding. wpe: word position embedding.
            wpe = nn.Embedding(config.block_size, config.n_embd), # wpe: trained position embedding.
            drop = nn.Dropout(config.dropout),
            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
            ln_f = LayerNorm(config.n_embd, bias=config.bias),
        ))
        # last layer. 
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        
        # with weight tying when using torch.compile() some warnings get generated:
        # "UserWarning: functional_call was passed multiple values for tied weights.
        # This behavior is deprecated and will be an error in future versions"
        # not 100% sure what this is, so far seems to be harmless. TODO investigate
                
        # lm_head: re-use the embedding table, computes the inner-product of inputs with each embedding.
        # the embedding layer (vocab_size -> n_emb )vs the last projection layer (n_embed -> vocab )
        self.transformer.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying

        # apply recursively for every subm-module
        self.apply(self._init_weights) 
        # adhoc: apply special scaled init to the residual projections, per GPT-2 paper
        for pn, p in self.named_parameters():
            if pn.endswith('c_proj.weight'):
                torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))

        # report number of parameters
        print("number of parameters: %.2fM" % (self.get_num_params()/1e6,))

    def get_num_params(self, non_embedding=True):
        """
        Return the number of parameters in the model.
        For non-embedding count (default), the position embeddings get subtracted.
        The token embeddings would too, except due to the parameter sharing these
        params are actually used as weights in the final layer, so we include them.
        """
        n_params = sum(p.numel() for p in self.parameters()) # numel() for total param.
        if non_embedding:
            n_params -= self.transformer.wpe.weight.numel() # remove wpe: shared param
        return n_params

    def _init_weights(self, module):
        # only linear layer and embedding?
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        # target is the same shape with idx. Each token has a target, and a loss
        device = idx.device # idx is a tensor.
        b, t = idx.size() # batch size, seq len, (seq of tokens)
        # seq length can be less than context window size.
        assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
        pos = torch.arange(0, t, dtype=torch.long, device=device) # shape (t)

        # forward the GPT model itself
        # wte: embedding layer, lookup according to input size [b,t], output [b,t,n_embd]
        tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)
        # wpe: look up [1,t] into [t, n_emd]
        pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd)
        # dropout at input
        x = self.transformer.drop(tok_emb + pos_emb) 
        # shape [B, T, C]
        for block in self.transformer.h:
            x = block(x)
        # layernorm after the output.
        x = self.transformer.ln_f(x)

        if targets is not None:
            # if we are given some desired targets also calculate the loss
            # shape [B, T, V] vocab_size, T sequence size.
            logits = self.lm_head(x)

            # logits.shape [64, 256, 65]
            # targets.shape [64, 256], toy vocabulary of 65, target.max()=64
            # targets.view(-1).shape: 16384, -1 means all element.s
            # logits.size(-1): 65
            # logits.view(-1, 65): shape [16384, 65], the two -1 have different meaning, size(-1) get the last dim, view(-1) get all elements
            # cross_entropy([16384, 65] [16384], ignore_index (the padding index)
            # IGNORE_IDX for computing cross_entropy
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=PAD_IDX)
        else:
            # inference-time mini-optimization: only forward the lm_head on the very last position
            logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim
            loss = None

        return logits, loss

    def crop_block_size(self, block_size):
        # model surgery to decrease the block size if necessary
        # e.g. we may load the GPT2 pretrained model checkpoint (block size 1024)
        # but want to use a smaller block size for some smaller, simpler model
        assert block_size <= self.config.block_size
        self.config.block_size = block_size
        self.transformer.wpe.weight = nn.Parameter(self.transformer.wpe.weight[:block_size])
        for block in self.transformer.h:
            if hasattr(block.attn, 'bias'): # only attention bias (the casual mask)
                block.attn.bias = block.attn.bias[:,:,:block_size,:block_size]

    def configure_optimizers(self, weight_decay, learning_rate, betas):
        # adam, learning rate, beta, weight_decay
        # start with all of the candidate parameters
        param_dict = {pn: p for pn, p in self.named_parameters()}
        # filter out those that do not require grad
        param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}
        
        # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.
        # i.e. all weight tensors in matmuls + embeddings decay,
        decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
        # all biases and layernorms don't.
        nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
        optim_groups = [
            {'params': decay_params, 'weight_decay': weight_decay},
            {'params': nodecay_params, 'weight_decay': 0.0}
        ]
        num_decay_params = sum(p.numel() for p in decay_params)
        num_nodecay_params = sum(p.numel() for p in nodecay_params)
        print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters")
        print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters")
        optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas)
        return optimizer

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        """
        Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete
        the sequence max_new_tokens times, feeding the predictions back into the model each time.
        Most likely you'll want to make sure to be in model.eval() mode of operation for this.

        idx: shape (b,t), contains the prompt.
        first iteration: computes all k-v cache
        follow iterations: re-use k-v cache, auto-regressive generating. 
        """
        for _ in range(max_new_tokens):
            # if the sequence context is growing too long we must crop it at block_size
            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:] # idx[:, -blocksize:], keep the last blocksize columsn
            # forward the model to get the logits for the index in the sequence
            logits, _ = self(idx_cond)
            # pluck the logits at the final step and scale by desired temperature
            # logit shape [B, T, V], use the last one in sequence.
            # divided by temperate > 1, reduces diversity , divied by temperatre < 1, thorugh linear scale in raw domain, nonline in exp domain.
            logits = logits[:, -1, :] / temperature 
            # optionally crop the logits to only the top k options
            if top_k is not None:
                # import ipdb;ipdb.set_trace()
                # v,_ the same shape with logits, v: values, _, index
                v, _ = torch.topk(logits, min(top_k, logits.size(-1))) # topk, return topk alone one dimension, [N, N], become [N,k]
                # logits [1,65], v[:, [-1]] scalar, 
                # logits < v[:, [-1]]
                # logits [b, t], v: [b, k], v[:, [-1]], [b,1] represented kth largest, 
                # logits < v[:, [-1]] broadcast [b, 1] into [b, t], then perform element-wise comparision .
                logits[logits < v[:, [-1]]] = -float('Inf') # mask out values below topk. 
            # apply softmax to convert logits to (normalized) probabilities
            probs = F.softmax(logits, dim=-1)            
            # opt 1: greedy decoding, choose the most probable. 
            # opt 2: top-K followed by multinomial sampling, no beam search, 
            idx_next = torch.multinomial(probs, num_samples=1)
            # append sampled index to the running sequence and continue
            idx = torch.cat((idx, idx_next), dim=1)

        return idx

 # %% Load pre-trained GTP

 from transformers import GPT2LMHeadModel
 # n_layer, n_head and n_embd are determined from model_type
 model_type = 'gpt2' # hardcoded 
 config_args = {
    'gpt2':  dict(n_layer=12, n_head=12, n_embd=768),  # 124M params
 }[model_type]
 print("forcing vocab_size=50257, block_size=1024, bias=True")
 config_args['vocab_size'] = 50257 # always 50257 for GPT model checkpoints
 config_args['block_size'] = 1024 # always 1024 for GPT model checkpoints
 config_args['bias'] = True # always True for GPT model checkpoints
 # create a from-scratch initialized minGPT model
 config = GPTConfig(**config_args)
 model = GPT(config)
 sd = model.state_dict()
 sd_keys = sd.keys()
 sd_keys = [k for k in sd_keys if not k.endswith('.attn.bias')] # discard this mask / buffer, not a param

 # init a huggingface/transformers model
 model_hf = GPT2LMHeadModel.from_pretrained(model_type)
 sd_hf = model_hf.state_dict()

 # copy while ensuring all of the parameters are aligned and match in names and shapes
 sd_keys_hf = sd_hf.keys()
 sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.masked_bias')] # ignore these, just a buffer
 sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.bias')] # same, just the mask (buffer)
 transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']
 # basically the openai checkpoints use a "Conv1D" module, but we only want to use a vanilla Linear
 # this means that we have to transpose these weights when we import them
 assert len(sd_keys_hf) == len(sd_keys), f"mismatched keys: {len(sd_keys_hf)} != {len(sd_keys)}"
 for k in sd_keys_hf:
    if any(k.endswith(w) for w in transposed):
        # special treatment for the Conv1D weights we need to transpose
        assert sd_hf[k].shape[::-1] == sd[k].shape #shape[::-1] reversed.
        with torch.no_grad():
            sd[k].copy_(sd_hf[k].t()) # copy_: shallow coyh, share memory
    else:
        # vanilla copy over the other parameters
        assert sd_hf[k].shape == sd[k].shape
        with torch.no_grad():
            sd[k].copy_(sd_hf[k])

 #%% validate GPT decoder
 device=torch.device("cuda") if torch.cuda.is_available() else torch.device('cpu')
 model.to(device)

 x,y = next(iter(train_dataloader))
 x,y = x.to(device), y.to(device)
 logits, loss = model.forward(x, y)
 print(x.shape)
 print(y.shape)
 print(logits.shape)
 print(loss)


 # %% fine-tune GPT on translation task. 
 learning_rate = 6e-4 # max learning rate
 max_epochs = 10
 weight_decay = 0.1
 beta1 = 0.9
 beta2 = 0.95 # beta1, beta2, hyperparameter in adam.
 num_train_batches = len(train_dataloader)
 optimizer = model.configure_optimizers(weight_decay, learning_rate, (beta1, beta2))

 from tqdm import tqdm
 for epoch in range(max_epochs):   
    # with context manager for tqdm.
    # with tqdm(total=num_train_batches, desc="Processing") as pbar:
    model.train()
    epoch_l = 0
    pbar = tqdm(total=num_train_batches, desc="training")  # start pbar
    for src,tgt in train_dataloader:
        src = src.to(device)
        tgt = tgt.to(device)

        optimizer.zero_grad() # start zero_grad
        logits,batch_loss = model.forward(src,tgt)
        epoch_l += batch_loss.item()
        batch_loss.backward()
        # if gradient_clip_val > 0:  # To be discussed later
        #     clip_gradients(gradient_clip_val, model)
        optimizer.step()
        pbar.update(1)
        pbar.set_postfix({"loss": batch_loss}) # Update average loss
    pbar.close() # manually close
No results found