temp_code.py

"""
import torch
import torch.nn as nn
from model.transformer import build_transformer
from utils.tokenizer import get_tokenizer
from data.dataloader import get_dataloaders

# ===================== Main Training Script ===================== #

def main():
  # ---------- Configuration ----------
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  src_vocab_size = 10000
  tgt_vocab_size = 10000
  src_seq_len = 50
  tgt_seq_len = 50
  d_model = 512
  n_layers = 6
  n_heads = 8
  dropout = 0.1
  d_ff = 2048
  epochs = 5
  batch_size = 32
  lr = 1e-4

  # ---------- Data ----------
  tokenizer = get_tokenizer()  # Load your tokenizer (custom or HuggingFace)
  train_loader, val_loader = get_dataloaders(tokenizer, batch_size)

  # ---------- Model ----------
  model = build_transformer(
    src_vocab_size, tgt_vocab_size,
    src_seq_len, tgt_seq_len,
    d_model, n_layers, n_heads,
    dropout, d_ff
  ).to(device)

  # ---------- Loss & Optimizer ----------
  criterion = nn.CrossEntropyLoss(ignore_index=0)  # Assuming <PAD> is 0
  optimizer = torch.optim.Adam(model.parameters(), lr=lr)

  # ---------- Training Loop ----------
  for epoch in range(epochs):
    model.train()
    total_loss = 0

    for src, tgt_input, tgt_output, src_mask, tgt_mask in train_loader:
      src = src.to(device)
      tgt_input = tgt_input.to(device)
      tgt_output = tgt_output.to(device)
      src_mask = src_mask.to(device)
      tgt_mask = tgt_mask.to(device)

      encoder_output = model.encode(src, src_mask)
      decoder_output = model.decode(encoder_output, src_mask, tgt_input, tgt_mask)
      logits = model.project(decoder_output)

      optimizer.zero_grad()
      loss = criterion(logits.view(-1, logits.size(-1)), tgt_output.view(-1))
      loss.backward()
      optimizer.step()
      total_loss += loss.item()

    avg_loss = total_loss / len(train_loader)
    print(f"Epoch [{epoch+1}/{epochs}], Loss: {avg_loss:.4f}")

if __name__ == "__main__":
  main()

"""

"""
import torch
from model.transformer import build_transformer
from utils.tokenizer import get_tokenizer
from data.dataloader import get_dataloaders

# ===================== Evaluation Script ===================== #

def evaluate():
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

  # ---------- Load Model ----------
  model = build_transformer(
    src_vocab_size=10000,
    tgt_vocab_size=10000,
    src_seq=50,
    tgt_seq=50,
    d_model=512,
    N=6,
    h=8,
    dropout=0.1,
    d_ff=2048
  )
  model.load_state_dict(torch.load("checkpoints/model.pth"))
  model.eval().to(device)

  # ---------- Load Tokenizer & Data ----------
  tokenizer = get_tokenizer()
  _, val_loader = get_dataloaders(tokenizer, batch_size=1)

  # ---------- Evaluation ----------
  with torch.no_grad():
    for src, tgt_input, tgt_output, src_mask, tgt_mask in val_loader:
      src = src.to(device)
      tgt_input = tgt_input.to(device)
      src_mask = src_mask.to(device)
      tgt_mask = tgt_mask.to(device)

      encoder_output = model.encode(src, src_mask)
      decoder_output = model.decode(encoder_output, src_mask, tgt_input, tgt_mask)
      output = model.project(decoder_output)

      predicted = torch.argmax(output, dim=-1)
      print("Predicted:", tokenizer.decode(predicted[0]))
      print("Target   :", tokenizer.decode(tgt_output[0]))
      break  # Only show one example

if __name__ == "__main__":
  evaluate()

"""

"""
import logging
import sys
from datetime import datetime


def get_logger(name: str, log_level=logging.INFO):
  logger = logging.getLogger(name)
  logger.setLevel(log_level)

  if not logger.handlers:
    # Create formatter
    formatter = logging.Formatter("[%(asctime)s] %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S")

    # Console handler
    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setFormatter(formatter)

    # Optional: file handler
    file_handler = logging.FileHandler(f"logs/{name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log")
    file_handler.setFormatter(formatter)

    # Add handlers
    logger.addHandler(console_handler)
    logger.addHandler(file_handler)

  return logger

"""
import os

checkpoint_dir = "checkpoints"
print("Available checkpoints:")
print(os.listdir(checkpoint_dir))