chore: improve comments and remove unnecessary blank lines

src/prediction.py

@@ -6,8 +6,8 @@ from services.model import load_model, get_device
 # Initialize tokenizer
 _tokenizer = tiktoken.get_encoding("cl100k_base")
 
+# Generate an imaginary word and its definition from three input words.
 def generate_word(words, model, vocab, inv_vocab, max_length=64):
-  """Generate an imaginary word and its definition from three input words."""
   device = get_device()
   
   # Tokenize input words

src/services/model.py

@@ -13,14 +13,8 @@ _VOCAB_PATH = os.path.join(_VOCAB_DIR, "vocab.json")
 # Internal device selection
 _DEVICE = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
 
+# Save the model state and vocabulary to disk.
 def save_model(model, vocab):
-  """
-  Save the model state and vocabulary to disk.
-  
-  Args:
-    model: The trained transformer model
-    vocab: The vocabulary dictionary
-  """
   # Create necessary directories
   os.makedirs(_MODEL_DIR, exist_ok=True)
   os.makedirs(_VOCAB_DIR, exist_ok=True)
@@ -35,13 +29,8 @@ def save_model(model, vocab):
   print(f"Model saved to {_MODEL_PATH}")
   print(f"Vocabulary saved to {_VOCAB_PATH}")
 
+# Load the model and its vocabulary from disk.
 def load_model():
-  """
-  Load the model and its vocabulary from disk.
-  
-  Returns:
-    tuple: (model, vocab, inv_vocab)
-  """
   # Load vocabulary
   with open(_VOCAB_PATH, "r", encoding="utf-8") as f:
     vocab = json.load(f)
@@ -54,11 +43,6 @@ def load_model():
   
   return model, vocab, inv_vocab
 
+# Get the device being used for model operations.
 def get_device():
-  """
-  Get the device being used for model operations.
-  
-  Returns:
-    torch.device: The device being used
-  """
   return _DEVICE 

src/services/tokenizer.py

@@ -4,8 +4,8 @@ from constants.tokens import special_tokens
 # Private tokenizer instance (internal use only)
 _tokenizer = tiktoken.get_encoding("cl100k_base")
 
+# Tokenize keys and values using the internal tokenizer.
 def tokenize_dataset(data):
-  """Tokenize keys and values using the internal tokenizer."""
   inputs = []
   outputs = []
   for key, value in data.items():
@@ -15,8 +15,8 @@ def tokenize_dataset(data):
     outputs.append(out_tokens)
   return inputs, outputs
 
+# Build vocabulary mapping from token IDs and add special tokens.
 def build_vocab(inputs, outputs):
-  """Build vocabulary mapping from token IDs and add special tokens."""
   offset = len(special_tokens)
   all_ids = set(tok for seq in inputs + outputs for tok in seq)
   vocab = {tok: i + offset for i, tok in enumerate(sorted(all_ids))}

src/services/transformer.py

@@ -1,38 +1,28 @@
 import torch
 import torch.nn as nn
-# import torch.nn.functional as F
 from constants.tokens import PAD_ID
 
 class TinyTransformer(nn.Module):
   def __init__(self, vocab_size, d_model=256, nhead=4, num_layers=2, dim_feedforward=512, dropout=0.1):
     super().__init__()
-    # self.pad_token_id = pad_token_id
-
     self.embedding = nn.Embedding(vocab_size, d_model, padding_idx=PAD_ID)
     self.pos_encoder = PositionalEncoding(d_model, dropout)
-
     encoder_layer = nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout, batch_first=True)
     self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
-
     decoder_layer = nn.TransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout, batch_first=True)
     self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
-
     self.out = nn.Linear(d_model, vocab_size)
 
+  # Keep tensors in batch-first format
   def forward(self, src, tgt):
-    # Keep tensors in batch-first format
     tgt_mask = nn.Transformer.generate_square_subsequent_mask(tgt.size(1)).to(src.device).bool()
-
     src_emb = self.pos_encoder(self.embedding(src))
     tgt_emb = self.pos_encoder(self.embedding(tgt))
-
     # Create padding masks
     src_padding_mask = (src == PAD_ID).bool()
     tgt_padding_mask = (tgt == PAD_ID).bool()
-
     memory = self.encoder(src_emb, src_key_padding_mask=src_padding_mask)
     output = self.decoder(tgt_emb, memory, tgt_mask=tgt_mask, tgt_key_padding_mask=tgt_padding_mask)
-
     return self.out(output)  # (batch, seq_len, vocab)
 
   def generate_src_mask(self, size):
@@ -42,7 +32,6 @@ class PositionalEncoding(nn.Module):
   def __init__(self, d_model, dropout=0.1, max_len=512):
     super().__init__()
     self.dropout = nn.Dropout(p=dropout)
-
     position = torch.arange(0, max_len).unsqueeze(1)
     div_term = torch.exp(
       torch.arange(0, d_model, 2) * (-torch.log(torch.tensor(10000.0)) / d_model)
@@ -50,7 +39,6 @@ class PositionalEncoding(nn.Module):
     pe = torch.zeros(max_len, d_model)
     pe[:, 0::2] = torch.sin(position * div_term)  # even indices
     pe[:, 1::2] = torch.cos(position * div_term)  # odd indices
-
     self.register_buffer('pe', pe.unsqueeze(0))
 
   def forward(self, x):