simonbb 0.0.1__tar.gz

simonbb-0.0.1/LICENSE

@@ -0,0 +1,19 @@
+Copyright (c) 2018 The Python Packaging Authority
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

simonbb-0.0.1/PKG-INFO

@@ -0,0 +1,16 @@
+Metadata-Version: 2.4
+Name: simonbb
+Version: 0.0.1
+Summary: A small package for teaching deep learning
+Project-URL: Homepage, https://simonbb.com
+Project-URL: Issues, https://simonbb.com/issues
+Author-email: Huaxia Rui <huaxia.rui@simon.rochester.edu>
+License-Expression: MIT
+License-File: LICENSE
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+
+# What's this package for?
+Anyone who teaches deep learning using PyTorch but doesn't want to bother students with too much nitty-gritty details.

simonbb-0.0.1/README.md

@@ -0,0 +1,2 @@
+# What's this package for?
+Anyone who teaches deep learning using PyTorch but doesn't want to bother students with too much nitty-gritty details.

simonbb-0.0.1/pyproject.toml

@@ -0,0 +1,23 @@
+[build-system]
+requires = ["hatchling >= 1.26"]
+build-backend = "hatchling.build"
+
+[project]
+name = "simonbb"
+version = "0.0.1"
+authors = [
+  { name="Huaxia Rui", email="huaxia.rui@simon.rochester.edu" },
+]
+description = "A small package for teaching deep learning"
+readme = "README.md"
+requires-python = ">=3.9"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Operating System :: OS Independent",
+]
+license = "MIT"
+license-files = ["LICEN[CS]E*"]
+
+[project.urls]
+Homepage = "https://simonbb.com"
+Issues = "https://simonbb.com/issues"

simonbb-0.0.1/src/simonbb/bert.py

@@ -0,0 +1,39 @@
+import re
+import torch
+import torch.nn as nn
+from transformers import BertTokenizer, BertModel, BertConfig
+
+#%%
+
+tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', do_lower_case=True)
+       
+class BertClassifier( nn.Module ):
+    def __init__(self, freeze_bert=False):
+        super().__init__()
+        config = BertConfig( max_position_embeddings=1024 ) # Instantiate BERT model, default hidden_size=767, i.e., dimensionality of the encoder layers and the pooler layer
+        BertModel( config )
+        self.bert = BertModel.from_pretrained('bert-base-uncased')
+        D_in, H, D_out = 768, 50, 2 # D_in is hidden_size of BERT, H is the hidden size of our classifier, and D_out is the number of classes
+        self.classifier = nn.Sequential( nn.Linear(D_in, H), nn.ReLU(), nn.Linear(H, D_out) )  # Instantiate a one-layer feed-forward classifier.  nn.Dropout(0.5),
+        if freeze_bert:
+            for param in self.bert.parameters():
+                param.requires_grad = False # if True, PyTorch will track all tensors that have param as an ancestor
+
+    def forward(self, x):
+        input_ids, attention_mask = x
+        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask) # calling an instance of nn.Module ends up calling forward() with the same arguments
+        last_hidden_state_cls = outputs[0][:, 0, :] # Extract the last hidden state of the token `[CLS]` for classification task
+        logits = self.classifier( last_hidden_state_cls )
+        return logits
+
+
+def BertPreprocessor( txt ):
+    input_ids = []
+    attention_masks = []
+    for s in txt:
+        s = re.sub(r'(@.*?)[\s]', ' ', s)
+        s = re.sub(r'\s+', ' ', s).strip()
+        encoded = tokenizer.encode_plus( text=s, add_special_tokens=True, max_length=128, padding='max_length', truncation=True, return_attention_mask=True )
+        input_ids.append( encoded.get('input_ids') ) # Add the outputs to the lists
+        attention_masks.append( encoded.get('attention_mask') )
+    return torch.tensor( input_ids ), torch.tensor( attention_masks )

simonbb-0.0.1/src/simonbb/gpt.py

@@ -0,0 +1,204 @@
+import numpy as np
+import torch
+import torch.nn as nn
+
+from .transformer import MultiHeadAttention, FeedForward  
+
+GPT2CONFIG = {"vocab_size": 50257, "context_length": 1024, "drop_rate": 0.0, "qkv_bias": True }
+GPT2SIZE = { "gpt2-small": {"emb_dim": 768, "n_layers": 12, "n_heads": 12, "size": "124M"},  # 621.83 MB
+             "gpt2-medium": {"emb_dim": 1024, "n_layers": 24, "n_heads": 16, "size":"355M"}, # 1549.58 MB
+             "gpt2-large": {"emb_dim": 1280, "n_layers": 36, "n_heads": 20, "size":"774M"},  # 3197.56 MB
+             "gpt2-xl": {"emb_dim": 1600, "n_layers": 48, "n_heads": 25, "size":"1558M"} }   # 6247.68 MB
+
+# Use "scale" and "shift" to match GPT2 implementation of LayerNorm. nn.LayerNorm() uses 'weight' and 'bias'
+class LayerNorm(nn.Module):
+    def __init__(self, emb_dim):
+        super().__init__()
+        self.eps = 1e-5
+        self.scale = nn.Parameter(torch.ones(emb_dim))
+        self.shift = nn.Parameter(torch.zeros(emb_dim))
+
+    def forward(self, x):
+        mean = x.mean(dim=-1, keepdim=True)
+        var = x.var(dim=-1, keepdim=True, unbiased=False)
+        norm_x = (x - mean) / torch.sqrt(var + self.eps)
+        return self.scale * norm_x + self.shift
+    
+class TransformerBlock(nn.Module): # Pre-LayerNorm, instead of Post-LayerNorm in the original paper
+    def __init__(self, cfg):
+        super().__init__()
+        self.att = MultiHeadAttention( embed_dim=cfg["emb_dim"], num_heads=cfg["n_heads"], seq_len=cfg["context_length"], dropout=cfg["drop_rate"], qkv_bias=cfg["qkv_bias"], is_causal=True )
+        self.ff = FeedForward( cfg["emb_dim"], 4*cfg["emb_dim"] )
+        self.norm1 = LayerNorm( cfg["emb_dim"] )
+        self.norm2 = LayerNorm( cfg["emb_dim"] )
+        self.drop_shortcut = nn.Dropout( cfg["drop_rate"] )
+
+    def forward(self, x):        
+        shortcut = x # Shortcut connection for attention block
+        x = self.norm1(x)
+        x = self.att(x,x,x)  # Shape [batch_size, num_tokens, emb_size]
+        x = self.drop_shortcut(x)
+        x = x + shortcut  # Add the original input back        
+        shortcut = x # Shortcut connection for feed forward block
+        x = self.norm2(x)
+        x = self.ff(x)
+        x = self.drop_shortcut(x)
+        x = x + shortcut  # Add the original input back
+        return x
+
+
+class GPTModel(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.tok_emb  = nn.Embedding( cfg["vocab_size"], cfg["emb_dim"] )
+        self.pos_emb  = nn.Embedding( cfg["context_length"], cfg["emb_dim"] )
+        self.drop_emb = nn.Dropout( cfg["drop_rate"] )
+        self.trf_blocks = nn.Sequential( *[TransformerBlock(cfg) for _ in range(cfg["n_layers"])] )       
+        self.final_norm = LayerNorm( cfg["emb_dim"] )
+        self.out_head   = nn.Linear( cfg["emb_dim"], cfg["vocab_size"], bias=False )
+
+    def forward(self, in_idx):
+        batch_size, seq_len = in_idx.shape
+        tok_embeds = self.tok_emb(in_idx)
+        pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
+        x = tok_embeds + pos_embeds  # Shape [batch_size, num_tokens, emb_size]
+        x = self.drop_emb(x)
+        x = self.trf_blocks(x)
+        x = self.final_norm(x)
+        logits = self.out_head(x) # Applying a PyTorch Linear layer to a rank-3 tensor performs a linear transformation only on the last dimension.
+        return logits
+
+
+#%% utility functions for training and inference
+
+def text_to_token_ids(text, tokenizer):
+    encoded = tokenizer.encode(text)
+    encoded_tensor = torch.tensor(encoded).unsqueeze(0)  # add batch dimension
+    return encoded_tensor
+
+def token_ids_to_text(token_ids, tokenizer):
+    flat = token_ids.squeeze(0)  # remove batch dimension
+    return tokenizer.decode(flat.tolist())
+
+# generate based on idx, context of shape (batch, seq_len) with each element a token ID
+def generate(model, idx, max_new_tokens, context_size, temperature=0.0, top_k=None, eos_id=None):    
+    for _ in range(max_new_tokens): 
+        idx_cond = idx[:, -context_size:] # crop current context to at most the last context_size number of token ids
+        with torch.no_grad():
+            logits = model(idx_cond) # (batch, seq_len, vocab_size)
+        logits = logits[:, -1, :] # By focusing on the last time step, the shape becomes (batch, vocab_size)
+        if top_k is not None: # keep only top_k values
+            top_logits, _ = torch.topk(logits, top_k) # returns top_k logits in descending order and their positions
+            min_val = top_logits[:, -1]
+            logits = torch.where(logits < min_val, torch.tensor(float('-inf')).to(logits.device), logits)
+        if temperature > 0.0: # apply temperature scaling
+            logits = logits / temperature
+            probs = torch.softmax(logits, dim=-1)  # logits.shape = (batch_size, vocab_size)            
+            idx_next = torch.multinomial(probs, num_samples=1) # idx_next.shape = (batch_size, 1)
+        else: # greedy selection of the idx of the vocab with the highest probability
+            idx_next = torch.argmax(logits, dim=-1, keepdim=True)  # idx_next.shape = (batch_size, 1)
+        if idx_next == eos_id:  # Stop generating early if end-of-sequence token is encountered and eos_id is specified
+            break
+        idx = torch.cat((idx, idx_next), dim=1) # idx.shape = (batch_size, seq_len+1) after appending id_next to idx
+    return idx
+
+#%% utility functions for loading pretrained GPT2 models
+
+def assign(left, right):
+    if left.shape != right.shape:
+        raise ValueError(f"Shape mismatch. Left: {left.shape}, Right: {right.shape}")
+    return torch.nn.Parameter(torch.tensor(right))
+
+def load_pretrained(gpt, params):
+    gpt.pos_emb.weight = assign(gpt.pos_emb.weight, params['wpe'])
+    gpt.tok_emb.weight = assign(gpt.tok_emb.weight, params['wte'])
+    for b in range(len(params["blocks"])):
+        q_w, k_w, v_w = np.split( (params["blocks"][b]["attn"]["c_attn"])["w"], 3, axis=-1 )
+        gpt.trf_blocks[b].att.W_q.weight = assign( gpt.trf_blocks[b].att.W_q.weight, q_w.T )
+        gpt.trf_blocks[b].att.W_k.weight = assign( gpt.trf_blocks[b].att.W_k.weight, k_w.T )
+        gpt.trf_blocks[b].att.W_v.weight = assign( gpt.trf_blocks[b].att.W_v.weight, v_w.T )
+        # load bias
+        q_b, k_b, v_b = np.split( (params["blocks"][b]["attn"]["c_attn"])["b"], 3, axis=-1 )
+        gpt.trf_blocks[b].att.W_q.bias = assign( gpt.trf_blocks[b].att.W_q.bias, q_b )
+        gpt.trf_blocks[b].att.W_k.bias = assign( gpt.trf_blocks[b].att.W_k.bias, k_b )
+        gpt.trf_blocks[b].att.W_v.bias = assign( gpt.trf_blocks[b].att.W_v.bias, v_b )
+        # load weights and bias for the output projection layer of mha
+        gpt.trf_blocks[b].att.out_proj.weight = assign( gpt.trf_blocks[b].att.out_proj.weight, params["blocks"][b]["attn"]["c_proj"]["w"].T )
+        gpt.trf_blocks[b].att.out_proj.bias = assign(   gpt.trf_blocks[b].att.out_proj.bias,   params["blocks"][b]["attn"]["c_proj"]["b"] )
+        # load weights and bias for feedforward layers
+        gpt.trf_blocks[b].ff.layers[0].weight = assign( gpt.trf_blocks[b].ff.layers[0].weight, params["blocks"][b]["mlp"]["c_fc"]["w"].T )
+        gpt.trf_blocks[b].ff.layers[0].bias = assign(   gpt.trf_blocks[b].ff.layers[0].bias,   params["blocks"][b]["mlp"]["c_fc"]["b"] )
+        gpt.trf_blocks[b].ff.layers[2].weight = assign( gpt.trf_blocks[b].ff.layers[2].weight, params["blocks"][b]["mlp"]["c_proj"]["w"].T )
+        gpt.trf_blocks[b].ff.layers[2].bias = assign(   gpt.trf_blocks[b].ff.layers[2].bias,   params["blocks"][b]["mlp"]["c_proj"]["b"] )
+        # load parameters for scale and shift of LayerNorm layers
+        gpt.trf_blocks[b].norm1.scale = assign( gpt.trf_blocks[b].norm1.scale, params["blocks"][b]["ln_1"]["g"] )
+        gpt.trf_blocks[b].norm1.shift = assign( gpt.trf_blocks[b].norm1.shift, params["blocks"][b]["ln_1"]["b"] )
+        gpt.trf_blocks[b].norm2.scale = assign( gpt.trf_blocks[b].norm2.scale, params["blocks"][b]["ln_2"]["g"] )
+        gpt.trf_blocks[b].norm2.shift = assign( gpt.trf_blocks[b].norm2.shift, params["blocks"][b]["ln_2"]["b"] )
+    gpt.final_norm.scale = assign( gpt.final_norm.scale, params["g"] )
+    gpt.final_norm.shift = assign( gpt.final_norm.shift, params["b"] )
+    # weight tying: the original GPT-2 model reused the token embedding weights in the output layer to reduce the total number of parameters.
+    gpt.out_head.weight  = assign( gpt.out_head.weight, params['wte'] )
+    
+#%%
+
+if __name__ == "__main__":
+    
+    choice = "gpt2-small"
+    cfg = GPT2CONFIG.copy()
+    cfg.update( GPT2SIZE[choice] )
+    cfg['qkv_bias'] = False
+    cfg['drop_rate'] = 0.1
+    
+    model = GPTModel( cfg )
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Total number of parameters: {total_params:,}")
+
+    print("Token embedding layer shape:", model.tok_emb.weight.shape) # nn.Embedding(num_embeddings, embedding_dim) has weight of shape (num_embeddings, embedding_dim)
+    print("Output layer shape:", model.out_head.weight.shape) # nn.Linear(in_features, out_features) has weight of shape (out_features,in_features)
+
+    total_params_gpt2 =  total_params - sum(p.numel() for p in model.out_head.parameters())
+    print(f"Number of trainable parameters considering weight tying: {total_params_gpt2:,}")
+
+    total_size_bytes = total_params * 4 # total size in bytes (assuming float32, 4 bytes per parameter)
+    print(f"Total size of the model: {total_size_bytes / (1024 * 1024):.2f} MB")
+    
+    torch.manual_seed(123)
+
+    x = torch.rand(2, 4, 768)  # Shape: [batch_size, num_tokens, emb_dim]
+    mha = MultiHeadAttention(d_in=768, d_out=768, context_length=1024, dropout=0.0, n_heads=12, is_causal=True)
+    print("context_vecs.shape:", mha(x).shape )
+
+    block = TransformerBlock( cfg )
+    print("Input shape:", x.shape)
+    print("Output shape:", block(x).shape)
+
+    import tiktoken
+    tokenizer = tiktoken.get_encoding("gpt2")
+    
+    batch = []
+    txt1 = "Every effort moves you"
+    txt2 = "Every day holds a"
+    batch.append( torch.tensor(tokenizer.encode(txt1)) )
+    batch.append( torch.tensor(tokenizer.encode(txt2)) )
+    batch = torch.stack(batch, dim=0)
+    print("Input batch:\n", batch)
+
+    out = model( batch )    
+    print("\nOutput shape:", out.shape)
+
+    start_context = "Hello, I am"
+    encoded = tokenizer.encode(start_context)
+    print("encoded:", encoded)
+
+    encoded_tensor = torch.tensor(encoded).unsqueeze(0)
+    print("encoded_tensor.shape:", encoded_tensor.shape)
+
+    model.eval() # disable dropout
+    out = generate( model=model, idx=encoded_tensor, max_new_tokens=6, context_size=cfg["context_length"] )
+
+    print("Output:", out)
+    print( tokenizer.decode(out.squeeze(0).tolist()) )
+    
+else: 
+    print(f'GPT imported from local file "{__name__}.py"')

simonbb-0.0.1/src/simonbb/transformer.py

@@ -0,0 +1,253 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import math
+
+#%% Positional Embedding
+
+def positional_encoding(seq_len, depth):
+    d = depth / 2
+    positions = np.arange(seq_len)[:, np.newaxis]  # shape: (seq_len, 1)
+    d = np.arange(d)[np.newaxis, :] / d  # shape: (1, d)
+    angle_rates = 1 / (10000 ** d)  # shape: (1, d)
+    angle_rads = positions * angle_rates  # shape: (seq_len, d)
+    pos_encoding = np.concatenate([np.sin(angle_rads), np.cos(angle_rads)], axis=-1)  # shape: (seq_len, 2*d)
+    return torch.tensor(pos_encoding, dtype=torch.float32)  # shape: (seq_len, depth)
+
+class PositionalEmbedding(nn.Module):
+    def __init__(self, vocab_size, d_emb, padding_idx=0, seq_len=1024):
+        super().__init__()
+        self.d_emb = d_emb  # d_emb is the dimension of positional embedding, shared by the encoder and the decoder
+        self.embedding = nn.Embedding(vocab_size, d_emb, padding_idx=padding_idx) # token embedding
+        self.register_buffer('pos_encoding', positional_encoding(seq_len=seq_len, depth=d_emb))  # position embedding, with maximum sequence length 2048
+        self._init_weights()
+    def _init_weights(self): # uniform initialization for embeddings, as in Keras        
+        nn.init.uniform_(self.embedding.weight, -0.05, 0.05) # PyTorch methods ending with an underscore makes changes in-place
+        self.embedding._fill_padding_idx_with_zero() # Keep padding as zeros
+    def forward(self, x):
+        seq_len = x.shape[1]
+        x = self.embedding(x)
+        x *= math.sqrt(self.d_emb)  # This factor sets the relative scale of the embedding and positonal encoding.
+        return x + self.pos_encoding[:seq_len, :].unsqueeze(0)  # use the position embedding matrix up to the seq_len position (i.e., row)
+
+#%% Attention and FeedForward
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, embed_dim, num_heads, dropout=0.0, seq_len=1024, qkv_bias=False, is_causal=False):
+        super().__init__()
+        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
+        self.d_model = embed_dim
+        self.n_heads = num_heads
+        self.d_k = embed_dim // num_heads
+        self.W_q = nn.Linear(embed_dim, embed_dim, bias=qkv_bias)
+        self.W_k = nn.Linear(embed_dim, embed_dim, bias=qkv_bias)
+        self.W_v = nn.Linear(embed_dim, embed_dim, bias=qkv_bias)
+        self.out_proj = nn.Linear(embed_dim, embed_dim)  # Linear layer to combine head outputs
+        self.dropout = nn.Dropout(dropout)
+        self.is_causal = is_causal
+        if self.is_causal: # fill 1 in upper triangular (excluding diagonal) part to indicate masking
+            self.register_buffer('mask', torch.triu(torch.ones(seq_len, seq_len), diagonal=1))
+        # register a tensor as part of the module's state, without it being a learnable parameter
+        self._init_weights()
+        
+    def _init_weights(self):        
+         for name, param in self.named_parameters():
+             if 'weight' in name:
+                 nn.init.xavier_uniform_(param) # Xavier/Glorot initialization (same as Keras default)
+             elif 'bias' in name:
+                 nn.init.zeros_(param)
+
+    def forward(self, q, k, v):
+        b, seq_len, d_in = q.shape
+        k = self.W_k(k) # Shape: (b, seq_len, d_model)
+        q = self.W_q(q)
+        v = self.W_v(v)
+        # Implicitly split the matrix by adding a `n_heads` dimension: (b, seq_len, d_model) -> (b, seq_len, n_heads, d_k)
+        k = k.view(b, seq_len, self.n_heads, self.d_k) # (b, seq_len, n_heads, d_k)
+        v = v.view(b, seq_len, self.n_heads, self.d_k) # (b, seq_len, n_heads, d_k)
+        q = q.view(b, seq_len, self.n_heads, self.d_k) # (b, seq_len, n_heads, d_k)    
+        k = k.transpose(1, 2) # (b, n_heads, seq_len, d_k)
+        q = q.transpose(1, 2) # (b, n_heads, seq_len, d_k)
+        v = v.transpose(1, 2) # (b, n_heads, seq_len, d_k)       
+        attn_scores = q @ k.transpose(2, 3) # (b, n_heads, seq_len, seq_len)
+        if self.is_causal:
+            mask_bool = self.mask.bool()[:seq_len, :seq_len] # truncate the mask to seq_len and convert to boolean         
+            attn_scores.masked_fill_(mask_bool, -torch.inf) # fill masked position with -torch.inf
+        attn_weights = torch.softmax(attn_scores / self.d_k**0.5, dim=-1) # (b, n_heads, seq_len, seq_len)
+        attn_weights = self.dropout(attn_weights) # (b, n_heads, seq_len, seq_len)
+        context_vec = (attn_weights @ v).transpose(1, 2) # (b, n_heads, seq_len, d_k), then (b, seq_len, n_heads, d_k) after transpose
+        context_vec = context_vec.contiguous().view(b, seq_len, self.d_model) # Combine heads: (b, seq_len, n_heads * d_k) where d_model = n_heads * d_k
+        context_vec = self.out_proj(context_vec) # optional projection
+        return context_vec
+
+    
+class FeedForward(nn.Module):
+    def __init__(self, d_emb, d_ff):
+        super().__init__()
+        self.layers = nn.Sequential(
+            nn.Linear(d_emb, d_ff),
+            nn.GELU(),
+            nn.Linear(d_ff, d_emb) )
+        self._init_weights()
+    def _init_weights(self): # Xavier/Glorot initialization (same as Keras default)
+        nn.init.xavier_uniform_(self.layers[0].weight)
+        nn.init.zeros_(self.layers[0].bias)
+        nn.init.xavier_uniform_(self.layers[2].weight)
+        nn.init.zeros_(self.layers[2].bias)
+    def forward(self, x):
+        return self.layers(x)
+    
+#%% Transformer = Encoder + Decoder + FeedForward
+
+
+class EncoderLayer(nn.Module):
+    def __init__(self, *, d_emb, n_heads, d_ff, dropout=0.1, prenorm=False): # default: post layernorm, as in the original paper
+        super().__init__()
+        self.gsa = MultiHeadAttention( embed_dim=d_emb, num_heads=n_heads, dropout=dropout )
+        self.ff = FeedForward(d_emb, d_ff)
+        self.norm1 = nn.LayerNorm( d_emb )
+        self.norm2 = nn.LayerNorm( d_emb )
+        self.drop_shortcut = nn.Dropout( dropout )
+        self.prenorm = prenorm
+    def forward(self, x):
+        if self.prenorm:
+            shortcut = x # Shortcut connection for attention block
+            x = self.norm1(x)
+        else:
+            x = self.norm1(x)
+            shortcut = x
+        x = self.gsa(x,x,x)  # Shape [batch_size, num_tokens, emb_size]
+        x = x + shortcut  # Add the original input back
+        if self.prenorm:
+            shortcut = x # Shortcut connection for feed forward block
+            x = self.norm2(x)
+        else:
+            x = self.norm2(x)
+            shortcut = x
+        x = self.ff(x)
+        x = self.drop_shortcut(x)
+        return x + shortcut  # Add the original input back
+
+
+class Encoder(nn.Module):
+    def __init__(self, *, n_layers, d_emb, n_heads, d_ff, dropout=0.1, prenorm=False):
+        super().__init__()
+        self.d_emb = d_emb
+        self.n_layers = n_layers
+        self.enc_layers = nn.ModuleList( [EncoderLayer(d_emb=d_emb, n_heads=n_heads, d_ff=d_ff, prenorm=prenorm) for _ in range(n_layers)] )
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):  # x is token embeddings shape: (batch, seq_len, d_emb)
+        x = self.dropout(x)
+        for i in range(self.n_layers):
+            x = self.enc_layers[i](x)
+        return x  # Shape: (batch_size, seq_len, d_emb)
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, *, d_emb, n_heads, d_ff, seq_len=1024, dropout=0.1, prenorm=False):
+        super().__init__()
+        self.csa = MultiHeadAttention( embed_dim=d_emb, num_heads=n_heads, dropout=dropout, seq_len=seq_len, is_causal=True)
+        self.ca = MultiHeadAttention( embed_dim=d_emb, num_heads=n_heads, dropout=dropout )
+        self.ff = FeedForward(d_emb, d_ff)
+        self.norm1 = nn.LayerNorm( d_emb )
+        self.norm2 = nn.LayerNorm( d_emb )
+        self.drop_shortcut = nn.Dropout( dropout )
+        self.prenorm = prenorm
+        
+    def forward(self, x, context):
+        if self.prenorm:
+            shortcut = x # Shortcut connection for attention block
+            x = self.norm1(x)
+        else:
+            x = self.norm1(x)
+            shortcut = x
+        x = self.csa(x,x,x)  # Shape [batch_size, num_tokens, emb_size]
+        x = self.ca(q=x, k=context, v=context)
+        x = x + shortcut  # Add the original input back
+        if self.prenorm:
+            shortcut = x # Shortcut connection for feed forward block
+            x = self.norm2(x)
+        else:
+            x = self.norm2(x)
+            shortcut = x
+        x = self.ff(x)
+        x = self.drop_shortcut(x)
+        return x + shortcut  # Add the original input back
+    
+
+class Decoder(nn.Module):
+    def __init__(self, *, n_layers, d_emb, n_heads, d_ff, seq_len=1024, dropout=0.1, prenorm=False):
+        super().__init__()
+        self.d_emb = d_emb
+        self.n_layers = n_layers
+        self.dropout = nn.Dropout(dropout)
+        self.dec_layers = nn.ModuleList( [DecoderLayer(d_emb=d_emb, n_heads=n_heads, d_ff=d_ff, seq_len=seq_len, dropout=dropout, prenorm=prenorm) for _ in range(n_layers)] )
+
+    def forward(self, x, context):  # x is of shape (batch, tgt_seq_len, d_emb); context is of shape (batch_size, context_len, d_emb)
+        x = self.dropout(x)
+        for i in range(self.n_layers):
+            x = self.dec_layers[i](x, context)
+        return x  # shape: (batch_size, tgt_seq_len, d_emb)
+
+
+class Transformer(nn.Module):
+    def __init__(self, *, n_layers, d_emb, n_heads, d_ff, src_vocab_size, tgt_vocab_size, seq_len=1024, dropout=0.1, prenorm=False ):
+        super().__init__()
+        self.src_pos_embedding = PositionalEmbedding(vocab_size=src_vocab_size, d_emb=d_emb, seq_len=seq_len)
+        self.tgt_pos_embedding = PositionalEmbedding(vocab_size=tgt_vocab_size, d_emb=d_emb, seq_len=seq_len)
+        self.encoder = Encoder(n_layers=n_layers, d_emb=d_emb, n_heads=n_heads, d_ff=d_ff, dropout=dropout, prenorm=prenorm)
+        self.decoder = Decoder(n_layers=n_layers, d_emb=d_emb, n_heads=n_heads, d_ff=d_ff, seq_len=seq_len, dropout=dropout, prenorm=prenorm)
+        self.final_layer = nn.Linear(d_emb, tgt_vocab_size)
+
+    def forward(self, x):
+        src, tgt = x
+        src_emb = self.src_pos_embedding(src)
+        tgt_emb = self.tgt_pos_embedding(tgt)
+        context = self.encoder(src_emb)  # (batch_size, context_len, d_emb)
+        x = self.decoder(tgt_emb, context)  # (batch_size, target_len, d_emb)
+        return self.final_layer(x)  # (batch_size, target_len, tgt_vocab_size)
+
+
+#%% Sanity Check: run from the parent folder of rui as "python -m rui.torch.transformer" since we used relative import
+if __name__ == "__main__":
+    
+    from ..utils import TextVectorizer
+    
+    seq_len = 10
+    max_vocab_size = 100
+    d_emb = 512
+    
+    vectorizer = TextVectorizer( max_tokens=max_vocab_size, output_mode="int" )
+    txts = ['how are you']
+    vectorizer.adapt( txts )
+    encoded = vectorizer(txts)
+   
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    
+    encoded_tensor = torch.tensor(encoded).to(device)
+
+    print(f"encoded_tensor: {encoded_tensor}")
+    
+    x = PositionalEmbedding(vocab_size=max_vocab_size, d_emb=d_emb, seq_len=seq_len).to(device)(encoded_tensor)
+    
+    print(f"input shape: {x.shape}")
+    n_heads = 8  
+    mha = MultiHeadAttention( d_emb, n_heads).to(device)
+    print(f"MultiHeadAttention output shape: {mha(x,x,x).shape}")
+
+    d_ff = 2048
+    encoder_layer = EncoderLayer(d_emb=d_emb, n_heads=n_heads, d_ff=d_ff).to(device)
+    decoder_layer = DecoderLayer(d_emb=d_emb, n_heads=n_heads, d_ff=d_ff).to(device)
+    print(f"EncoderLayer output shape: {encoder_layer(x).shape}")
+
+    n_layers = 6
+    encoder = Encoder(n_layers=n_layers, d_emb=d_emb, n_heads=n_heads, d_ff=d_ff).to(device)
+    decoder = Decoder(n_layers=n_layers, d_emb=d_emb, n_heads=n_heads, d_ff=d_ff).to(device)
+    encoder.eval()
+    decoder.eval()
+    with torch.no_grad():
+        print(f"Encoder output shape: {encoder(x).shape}")
+        print(f"Decoder output shape: {decoder(x=x, context=x).shape}")
+else:
+    print(f'Transformer imported from local file "{__name__}.py"')

simonbb-0.0.1/src/simonbb/utils.py

@@ -0,0 +1,486 @@
+"""
+A collection of utility functions for teaching deep learning
+@author: huaxia
+"""
+
+import os
+import time
+import numpy as np
+import copy
+import re
+import pickle
+import json
+from collections import Counter
+import urllib.request
+
+import matplotlib.pyplot as plt
+
+import torch
+import torch.nn as nn
+
+#%%
+
+class TextVectorizer:
+    def __init__(self, max_tokens=20000, output_mode="int", output_sequence_length=None, ngrams=None, standardize=None):
+        self.max_tokens = max_tokens
+        self.output_mode = output_mode
+        self.output_sequence_length = output_sequence_length
+        self.ngrams = ngrams
+        self.standardize = standardize
+        self.vocab = None
+        self.word_to_idx = None
+        self.idx_to_word = None
+        self.idf_weights = None
+
+    def _tokenize(self, text):
+        if self.standardize:
+            text = self.standardize(text)
+        else:
+            text = text.lower()
+            text = re.sub(r'[^\w\s]', ' ', text)
+        tokens = text.split()
+        if self.ngrams and self.ngrams >= 2:
+            ngram_tokens = []
+            for n in range(1, self.ngrams + 1):
+                for i in range( len(tokens) - n + 1 ):
+                    ngram_tokens.append(' '.join(tokens[i:i + n]))
+            return ngram_tokens
+        return tokens
+
+    def adapt(self, texts):
+        """Build vocabulary from texts"""
+        token_counter = Counter()
+        doc_freq = Counter()  # for TF-IDF
+        num_docs = 0
+        print("Adapting TextVectorizer...")
+        for text in texts:
+            num_docs += 1
+            tokens = self._tokenize(text)
+            token_counter.update(tokens)
+            doc_freq.update(set(tokens))
+        # Reserve 0 for padding, 1 for unknown, hence minus 2 to account for them.
+        most_common = token_counter.most_common(self.max_tokens - 2 if self.max_tokens else None) # most_common(n=None) returns all elements in the counter if n is omitted or None.
+        self.vocab = ['<pad>', '<unk>'] + [word for word, _ in most_common] # 0 for padding, 1 for unknown
+        self.word_to_idx = {word: idx for idx, word in enumerate(self.vocab)}
+        self.idx_to_word = {idx: word for idx, word in enumerate(self.vocab)}
+        # Compute IDF weights for tf-idf mode
+        if self.output_mode == "tf_idf":
+            self.idf_weights = np.zeros(len(self.vocab))
+            for word, idx in self.word_to_idx.items():
+                df = doc_freq.get(word, 0) + 1  # add 1 for smoothing
+                self.idf_weights[idx] = np.log(num_docs / df) + 1
+        print(f"Vocabulary size: {len(self.vocab)}")
+
+    def __call__(self, texts):
+        if isinstance(texts, str):
+            texts = [texts]
+        if self.output_mode == "int":
+            return self._vectorize_int(texts)
+        elif self.output_mode == "multi_hot":
+            return self._vectorize_multi_hot(texts)
+        elif self.output_mode == "tf_idf":
+            return self._vectorize_tfidf(texts)
+
+    def _vectorize_int(self, texts):
+        batch_indices = []
+        for text in texts:
+            tokens = self._tokenize(text)
+            indices = [self.word_to_idx.get(t, 1) for t in tokens]  # 1 is <unk>
+            if self.output_sequence_length:
+                if len(indices) < self.output_sequence_length:
+                    indices = indices + [0] * (self.output_sequence_length - len(indices))
+                else:
+                    indices = indices[:self.output_sequence_length]
+            batch_indices.append(indices)
+        return np.array(batch_indices)
+
+    def _vectorize_multi_hot(self, texts):
+        result = np.zeros( (len(texts), len(self.vocab)), dtype=np.float32 )
+        for i, text in enumerate(texts):
+            tokens = self._tokenize(text)
+            for t in tokens:
+                idx = self.word_to_idx.get(t, 1)
+                if idx < len(self.vocab):
+                    result[i, idx] = 1
+        return result
+
+    def _vectorize_tfidf(self, texts):
+        result = np.zeros( (len(texts), len(self.vocab)), dtype=np.float32 )
+        for i, text in enumerate(texts):
+            tokens = self._tokenize(text)
+            tf = Counter(tokens)
+            for t, count in tf.items():
+                idx = self.word_to_idx.get(t, 1)
+                if idx < len(self.vocab):
+                    result[i, idx] = count * self.idf_weights[idx]
+        return result
+
+    def get_vocabulary(self):
+        return self.vocab
+    
+    def save(self, path):
+        vec = {'max_tokens': self.max_tokens, 'output_sequence_length': self.output_sequence_length,
+               'vocab': self.vocab, 'word_to_idx': self.word_to_idx, 'idx_to_word': self.idx_to_word,
+               'has_standardizer': self.standardize is not None}
+        with open(path, 'wb') as f:
+            pickle.dump(vec, f)
+            
+    # transforms a method into a class method (i.e., bound to the class rather than an individual instance)
+    @classmethod
+    def load(cls, path, standardize=None):
+        data = pickle.load(path, weights_only=False)
+        vectorizer = cls( max_tokens=data['max_tokens'], output_sequence_length=data['output_sequence_length'], 
+                          standardize=standardize if data.get('has_standardize') else None )
+        vectorizer.vocab = data['vocab']
+        vectorizer.word_to_idx = data['word_to_idx']
+        vectorizer.idx_to_word = data['idx_to_word']
+        return vectorizer
+
+
+def query_ollama(prompt, model="llama3", url="http://localhost:11434/api/chat"):
+    # Create the data payload as a dictionary
+    data = {
+        "model": model,
+        "messages": [
+            {"role": "user", "content": prompt}
+        ],
+        "options": {     # Settings below are required for deterministic responses
+            "seed": 123,
+            "temperature": 0,
+            "num_ctx": 2048
+        }
+    }
+
+    # Convert the dictionary to a JSON formatted string and encode it to bytes
+    payload = json.dumps(data).encode("utf-8")
+
+    # Create a request object, setting the method to POST and adding necessary headers
+    request = urllib.request.Request(url, data=payload, method="POST")
+    request.add_header("Content-Type", "application/json")
+
+    # Send the request and capture the response
+    response_data = ""
+    with urllib.request.urlopen(request) as response:
+        # Read and decode the response
+        while True:
+            line = response.readline().decode("utf-8")
+            if not line:
+                break
+            response_json = json.loads(line)
+            response_data += response_json["message"]["content"]
+
+    return response_data
+
+#%%
+
+class ModelCheckpoint:
+    def __init__(self, filepath, monitor="val_loss", save_best_only=True, save_optimizer_state=False):
+        self.filepath = filepath
+        self.monitor = monitor
+        self.save_best_only = save_best_only
+        self.save_optimizer_state = save_optimizer_state
+        os.makedirs(os.path.dirname(filepath), exist_ok=True)
+
+    def on_best(self, model, optimizer):
+        if self.save_optimizer_state:
+            torch.save( {"model_state_dict": model.state_dict(), "optimizer_state_dict": optimizer.state_dict()}, self.filepath)
+        else:
+            torch.save(model.state_dict(), self.filepath)
+
+
+def plotEpoch(history, metric="average_loss"):
+    epochs = range(len(history["average_train_loss"]))
+    if metric == "accuracy":
+        plt.plot(epochs, history["train_accuracy"], "k", label="Training acc")
+        plt.plot(epochs, history["val_accuracy"], "b", label="Validation acc")
+        plt.axvline( x=int(np.argmax(history["val_accuracy"])), color="0.5" )
+        plt.title("Training and Validation Accuracy")
+        plt.legend()
+    elif metric == "total_loss":
+        plt.plot(epochs, history["total_train_loss"], "k", label="Training loss")
+        plt.plot(epochs, history["total_val_loss"], "b", label="Validation loss")
+        plt.axvline(x=int(np.argmin(history["total_val_loss"])), color="0.5")
+        plt.title("Total Training and Validation Loss")
+        plt.legend()
+    elif metric == "average_loss":
+        plt.plot(epochs, history["average_train_loss"], "k", label="Training loss")
+        plt.plot(epochs, history["average_val_loss"], "b", label="Validation loss")
+        plt.axvline(x=int(np.argmin(history["average_val_loss"])), color="0.5")
+        plt.title("Average Training and Validation Loss")
+        plt.legend()
+    plt.tight_layout()
+    plt.show()
+    plt.close()    
+
+
+def stats_per_channel(loader):
+    n_channels, _, _ = loader.dataset[0][0].size() # assuming each batch is a tuple of (x,y)
+    #total_samples = torch.ones( (1, n_channels) ) * len(loader.dataset)
+    total_means = torch.zeros( (1,n_channels) )
+    total_stds = torch.zeros( (1,n_channels) )
+    for images, _ in loader:        
+        flatten_per_channel = torch.flatten(images, start_dim=-2, end_dim=-1) # (N,C,H*W)
+        means = flatten_per_channel.mean(axis=2) # (N,C)
+        stds = flatten_per_channel.std(axis=2) # (N,C)
+        total_means += means.sum(axis=0) # (1, n_channels)
+        total_stds += stds.sum(axis=0) # (1, n_channels)
+    norm_mean = total_means / len(loader.dataset) # (1, n_channels)
+    norm_std = total_stds / len(loader.dataset) # (1, n_channels)
+    return norm_mean.squeeze(), norm_std.squeeze() # without squeezing, runtime error: not all boolean value of tensor with more than one value is ambiguous
+
+#%% Train Loop: tested on classification, regression, image segmentation, translation, gpt2
+
+def train( model, train_dataloader, val_dataloader, optimizer, scheduler=None, clip_grad=None, loss_fn=nn.CrossEntropyLoss(), callbacks=None, device=None, evaluation=True, n_epochs=1, n_batch_per_report=5, accuracy_fn=None, lr_history=None ):
+    if device is None:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device) # Models are moved in-place (the object itself is modified)
+    
+    batch = next(iter(train_dataloader))
+    y_shape = batch[-1].squeeze().shape
+    # y_shape is typically a scalar for scalar label, 1D tensor for sequence, and 2D for image segmentation task where the label is a 2D tensor representing an image
+    if len(y_shape)==1:
+        n_labels_per_instance = 1 # typical case
+    elif len(y_shape)==2: # (batch, seq_len) for sequence data
+        n_labels_per_instance = y_shape[1] 
+    elif len(y_shape)==3: # (batch, height, width) for image segmentation task
+        n_labels_per_instance = y_shape[1] * y_shape[2]
+    else:
+        print( f"Error: Targe tensors with {len(y_shape)} or more axes are currently not supported!" )
+
+    history = { "total_train_loss":[], "average_train_loss":[], "train_accuracy":[], "total_val_loss":[], "average_val_loss":[], "val_accuracy":[] }
+    best_val_loss = float("inf")
+    best_state = None
+
+    print(f'Start training on device {device}.\n')
+    for epoch in range(1, n_epochs + 1):
+        if isinstance(loss_fn, nn.MSELoss):
+            print(f"{'Epoch':^7} | {'Batch':^7} | {'Train Loss':^12} | {'Val Loss':^10} | {'Train MAE':^9} | {'Val MAE':^9} | {'Elapsed':^9}")
+        else:
+            print(f"{'Epoch':^7} | {'Batch':^7} | {'Train Loss':^12} | {'Val Loss':^10} | {'Train Acc':^9} | {'Val Acc':^9} | {'Elapsed':^9}")        
+        print("-" * 80)
+        
+        model.train()
+        t0_epoch, t0_batch = time.time(), time.time()
+        total_correct, total_loss, total_abs_err, weighted_accuracy, batch_loss, batch_count = 0, 0, 0, 0, 0, 0
+    
+        for batch in train_dataloader:
+            batch_count += 1
+            
+            if isinstance(loss_fn, nn.MSELoss): # regression
+                y = batch[-1].to(device)
+            else: # classification task
+                y = batch[-1].long().to(device) # convert to long (default integer type in PyTorch), required by nn.CrossEntropyLoss()
+                
+            if len(batch)>2: # example: batch is a triple (sequence, attention mask, label)                
+                x = tuple(t.to(device) for t in batch[:-1])
+            else:                
+                x = batch[0].to(device)
+                
+            optimizer.zero_grad(set_to_none=True)
+            
+            logits = model(x).squeeze() # squeeze for regression; otherwise, broadcasting amplifies loss and abs_err
+            
+            if len(y_shape)==1 or len(y_shape)==3: # scalar target or 2D image target
+                loss = loss_fn( logits, y )
+            elif len(y_shape)==2: # sequence target
+                loss = loss_fn( logits.flatten(0, 1), y.flatten() ) # pred: (batch * seq_len, vocab_size), label: (batch * seq_len,)
+            else:
+                print( f"Error: Targe tensors with {len(y_shape)} or more axes are currently not supported!" )
+                
+            batch_loss = loss.item() # item() automatically moves a single-element tensor's value to the CPU and returns it as a standard Python number.
+            total_loss += batch_loss * y.shape[0] # Default for CrossEntropyLoss is 'mean' which averages over each loss element in the batch
+            
+            if isinstance(loss_fn, nn.MSELoss):
+                abs_err = torch.sum( torch.abs(logits-y) ).item()
+                total_abs_err += abs_err
+                batch_accuracy = abs_err / y.shape[0]
+            elif accuracy_fn is None:
+                if len(y_shape)==1 or len(y_shape)==3: # scalar target or 2D image target
+                    preds = torch.argmax(logits, dim=1)
+                elif len(y_shape)==2:
+                    preds = torch.argmax(logits, dim=2) # sequence target
+                else:
+                    print( f"Error: Targe tensors with {len(y_shape)} or more axes are currently not supported!" )
+                batch_accuracy = (preds == y).cpu().numpy().mean() * 100
+                total_correct += (preds == y).sum().item() # item() automatically moves a single-element tensor's value to the CPU and returns it as a standard Python number.
+            else:
+                batch_accuracy = accuracy_fn( logits, y ).item() * 100
+                weighted_accuracy += batch_accuracy * len(y)
+                        
+            loss.backward()
+            
+            if clip_grad is not None: # gradient clipping prevents exploding gradients
+                torch.nn.utils.clip_grad_norm_(model.parameters(), clip_grad)
+
+            optimizer.step() # Update parameters
+            if scheduler is not None:
+                scheduler.step() # Update learning rate
+            if lr_history is not None:
+                lr_history.append( optimizer.param_groups[0]['lr'] )
+                        
+            if batch_count % n_batch_per_report == 0:
+                time_elapsed = time.time() - t0_batch
+                print( f"{epoch:^7} | {batch_count:^7} | {batch_loss:^12.6f} | {'-':^10} | {batch_accuracy:^9.2f} | {'-':^9} | {time_elapsed:^9.2f}")                           
+                
+            batch_loss = 0
+            t0_batch = time.time()
+
+        if n_batch_per_report < batch_count:
+            print("-" * 80)
+            
+        history["total_train_loss"].append( total_loss )
+        history['average_train_loss'].append( history["total_train_loss"][-1] / len(train_dataloader.dataset) )
+        
+        if isinstance(loss_fn, nn.MSELoss):
+            history["train_accuracy"].append( total_abs_err / len(train_dataloader.dataset) ) 
+        elif accuracy_fn is None: # scalar target (1), sequence prediction (2), image segmentation (3)
+            history["train_accuracy"].append( 100 * total_correct / ( len(train_dataloader.dataset)*n_labels_per_instance ) )
+        else:
+            history["train_accuracy"].append( weighted_accuracy / len(train_dataloader.dataset) ) # batch-size-weighted accuracy
+                
+        if evaluation == True:            
+            model.eval()
+            val_correct, val_loss, total_abs_err, weighted_accuracy = 0, 0, 0, 0
+            for batch in val_dataloader:
+                if isinstance(loss_fn, nn.MSELoss): # regression task
+                    y = batch[-1].to(device)
+                else: # classification task
+                    y = batch[-1].long().to(device) # convert to long (default integer type in PyTorch), required by nn.CrossEntropyLoss()                    
+                if len(batch)>2: # example: (sequence, attention mask)
+                    x = tuple(t.to(device) for t in batch[:-1])
+                else:
+                    x = batch[0].to(device)
+                    
+                with torch.no_grad(): # switch off autograd using the torch.no_grad() context manager
+                    logits = model( x ).squeeze() # squeeze for regression; otherwise, broadcasting amplifies loss and abs_err
+                    
+                if len(y_shape)==1 or len(y_shape)==3: # scalar target or 2D image target
+                    loss = loss_fn( logits, y )
+                elif len(y_shape)==2: # sequence target
+                    loss = loss_fn( logits.flatten(0, 1), y.flatten() ) # pred: (batch * seq_len, vocab_size), label: (batch * seq_len,)
+                else:
+                    print( f"Error: Targe tensors with {len(y_shape)} or more axes are currently not supported!" )
+                    
+                val_loss += loss.item() * len(y)  # Default for CrossEntropyLoss is 'mean' which averages over each loss element in the batch
+                
+                if isinstance(loss_fn, nn.MSELoss):
+                    total_abs_err += torch.sum( torch.abs(logits-y) ).item()
+                elif accuracy_fn is None:
+                    if len(y_shape)==1 or len(y_shape)==3: # scalar target or 2D image target
+                        preds = torch.argmax(logits, dim=1)
+                    elif len(y_shape)==2:
+                        preds = torch.argmax(logits, dim=2) # sequence target
+                    else:
+                        print( f"Error: Targe tensors with {len(y_shape)} or more axes are currently not supported!" )                    
+                    val_correct += (preds == y).cpu().numpy().sum().item()
+                else:
+                    batch_accuracy = accuracy_fn(logits, y).item()
+                    weighted_accuracy += batch_accuracy * len(y)
+                
+            history["total_val_loss"].append( val_loss )
+            history['average_val_loss'].append( history["total_val_loss"][-1] / len(val_dataloader.dataset) )
+            
+            if isinstance(loss_fn, nn.MSELoss):
+                history["val_accuracy"].append( total_abs_err / len(val_dataloader.dataset) ) 
+            elif accuracy_fn is None: # scalar target (1), sequence prediction (2), image segmentation (3)
+                history["val_accuracy"].append( 100 * val_correct / (len(val_dataloader.dataset)*n_labels_per_instance) )
+            else:
+                history["val_accuracy"].append( 100 * weighted_accuracy / len(val_dataloader.dataset) ) # batch-size-weighted accuracy
+            
+            time_elapsed = time.time() - t0_epoch    
+            print( f"{epoch:^7} | {'-':^7} | {history['average_train_loss'][-1]:^12.6f} | {history['average_val_loss'][-1]:^10.6f} | {history['train_accuracy'][-1]:^9.2f} | {history['val_accuracy'][-1]:^9.2f} | {time_elapsed:^9.2f}")
+            # ModelCheckpoint(save_best_only=True, monitor="val_loss")
+            if val_loss < best_val_loss:
+                best_val_loss = val_loss
+                best_state = copy.deepcopy( model.state_dict() )
+                optimizer_state = copy.deepcopy( optimizer.state_dict() )
+                if callbacks:
+                    for cb in callbacks:
+                        cb.on_best(model, optimizer)
+        else:
+            print( f"{epoch:^7} | {'-':^7} | {history['total_train_loss'][-1]:^12.6f} | {'-':^10} | {history['train_accuracy'][-1]:^9.2f} | {'-':^10} | {time_elapsed:^9.2f}")
+        print("-" * 80)
+    # restore best, similar to Keras's best-checkpoint behavior
+    if best_state is not None:
+        model.load_state_dict(best_state)
+        optimizer.load_state_dict(optimizer_state)
+    return history
+
+
+def evaluate(model, loader, device, loss_fn=nn.CrossEntropyLoss(), accuracy_fn=None):
+    model.eval()
+    batch = next(iter(loader))
+    y_shape = batch[-1].squeeze().shape
+    if len(y_shape)==1:
+        n_labels_per_instance = 1 # typical case
+    elif len(y_shape)==2: # (batch, seq_len) for sequence data
+        n_labels_per_instance = y_shape[1] 
+    elif len(y_shape)==3: # (batch, height, width) for image segmentation task
+        n_labels_per_instance = y_shape[1] * y_shape[2]
+    else:
+        print( f"Error: Targe tensors with {len(y_shape)} or more axes are currently not supported!" )
+        
+    val_correct, val_loss, total_abs_err, weighted_accuracy = 0,0,0,0
+    for batch in loader:
+        if isinstance(loss_fn, nn.MSELoss): # regression task
+            y = batch[-1].to(device)
+        else: # classification task
+            y = batch[-1].long().to(device) # convert to long (default integer type in PyTorch), required by nn.CrossEntropyLoss()
+        if len(batch)>2: # example: (sequence, attention mask)
+            x = tuple(t.to(device) for t in batch[:-1])
+        else:
+            x = batch[0].to(device)
+            
+        with torch.no_grad(): # switch off autograd using the torch.no_grad() context manager
+            logits = model(x).squeeze() # squeeze for regression; otherwise, broadcasting amplifies loss and abs_err
+            
+        if len(y_shape)==1 or len(y_shape)==3: # scalar target or 2D image target
+            loss = loss_fn( logits, y )
+        elif len(y_shape)==2: # sequence target
+            loss = loss_fn( logits.flatten(0, 1), y.flatten() ) # pred: (batch * seq_len, vocab_size), label: (batch * seq_len,)
+        else:
+            print( f"Error: Targe tensors with {len(y_shape)} or more axes are currently not supported!" )
+            
+        val_loss += loss.item() * y.shape[0]
+        
+        if isinstance(loss_fn, nn.MSELoss):
+            total_abs_err += torch.sum( torch.abs(logits-y) ).item()
+        elif accuracy_fn is None:
+            if len(y_shape)==1 or len(y_shape)==3: # scalar target or 2D image target
+                preds = torch.argmax(logits, dim=1)
+            elif len(y_shape)==2:
+                preds = torch.argmax(logits, dim=2) # sequence target: (batch, seq_len, vocab_size)
+            else:
+                print( f"Error: Targe tensors with {len(y_shape)} or more axes are currently not supported!" )              
+            val_correct += (preds == y).cpu().numpy().sum()
+        else:
+            batch_accuracy = accuracy_fn(logits, y).item()
+            weighted_accuracy += batch_accuracy * len(y)
+            
+    val_loss = val_loss / len(loader.dataset)
+    
+    if isinstance(loss_fn, nn.MSELoss):
+        val_accuracy = total_abs_err / len(loader.dataset)
+    elif accuracy_fn is None:
+        val_accuracy = val_correct/ (len(loader.dataset)**n_labels_per_instance)
+    else:
+        val_accuracy = 100 * weighted_accuracy / len(loader.dataset) # batch-size-weighted accuracy
+        
+    if torch.is_tensor(val_accuracy):
+        val_accuracy = val_accuracy.item() # extract the value of a single-item tensor 
+    return {'loss':val_loss, 'accuracy':val_accuracy}
+
+#%%
+
+if __name__ == "__main__":
+    
+    vectorizer = TextVectorizer( max_tokens=10, output_mode="int" )
+    txts = ['how are you']
+    vectorizer.adapt( txts )
+    encoded = vectorizer(txts)
+    print(encoded)
+    print(vectorizer.vocab)
+    print(vectorizer.word_to_idx)
+    print(vectorizer.idx_to_word)
+else:
+    print(f'Importing from local file "{__name__}.py"')