train.py

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
import re
import os
from tqdm import tqdm
import time

# --- Config ---
HIDDEN_SIZE = 128
BATCH_SIZE = 64
NUM_EPOCHS = 1
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
MODEL_DIR = "models"

# --- Data Utils ---
def load_words(filename):
    pattern = re.compile(r'^[a-zA-Z]+$')
    words = []
    with open(filename, 'r') as f:
        for line in f:
            w = line.strip()
            if pattern.match(w):
                words.append(w.lower())
    return words

def build_vocab(words):
    # Reserve 0=PAD, 1=<sos>, 2=<eos>
    chars = sorted(set(''.join(words)))
    char2idx = {'<pad>': 0, '<sos>': 1, '<eos>': 2}
    for i, c in enumerate(chars, start=3):
        char2idx[c] = i
    idx2char = {i: c for c, i in char2idx.items()}
    return char2idx, idx2char

# --- Obfuscator (Caesar) ---
def obfuscate_caesar(text, shift=3):
    result = ''
    for ch in text:
        if ch.isalpha():
            base = ord('A') if ch.isupper() else ord('a')
            result += chr((ord(ch) - base + shift) % 26 + base)
        else:
            result += ch
    return result

CURRENT_OBFUSCATOR = obfuscate_caesar

# --- Encoding helper ---
def encode(text, max_len, char2idx, add_tokens=False):
    seq = []
    if add_tokens:
        seq.append(char2idx['<sos>'])
    for ch in text:
        seq.append(char2idx.get(ch, 0))
    if add_tokens:
        seq.append(char2idx['<eos>'])
    # pad
    seq += [char2idx['<pad>']] * (max_len - len(seq))
    return seq

# --- Dataset ---
class StringDataset(Dataset):
    def __init__(self, words, obfuscator, char2idx, max_len):
        self.data = []
        for w in words:
            obf = obfuscator(w)
            src = encode(obf, max_len, char2idx, add_tokens=True)
            tgt = encode(w,   max_len, char2idx, add_tokens=True)
            self.data.append((src, tgt))
    def __len__(self):
        return len(self.data)
    def __getitem__(self, idx):
        return (
            torch.tensor(self.data[idx][0], dtype=torch.long),
            torch.tensor(self.data[idx][1], dtype=torch.long)
        )

# --- Positional Encoding ---
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=512):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2, dtype=torch.float) *
                             (-torch.log(torch.tensor(10000.0)) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe.unsqueeze(0))  # shape (1, max_len, d_model)

    def forward(self, x):
        # x: (batch, seq_len, d_model)
        x = x + self.pe[:, :x.size(1), :]
        return x

# --- Transformer Model ---
class TransformerModel(nn.Module):
    def __init__(self, vocab_size, hidden_size, max_len):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, hidden_size, padding_idx=0)
        self.pos_enc   = PositionalEncoding(hidden_size, max_len)
        self.transformer = nn.Transformer(
            d_model=hidden_size,
            nhead=4,
            num_encoder_layers=3,
            num_decoder_layers=3,
            dim_feedforward=256,
            dropout=0.1,
            batch_first=True
        )
        self.out = nn.Linear(hidden_size, vocab_size)

    def generate_square_subsequent_mask(self, sz):
        # mask out future positions
        return torch.triu(torch.full((sz, sz), float('-inf')), diagonal=1).to(DEVICE)

    def forward(self, src, tgt):
        # src, tgt: (batch, seq_len)
        src_emb = self.pos_enc(self.embedding(src))
        tgt_emb = self.pos_enc(self.embedding(tgt))
        tgt_mask = self.generate_square_subsequent_mask(tgt.size(1))
        out = self.transformer(src_emb, tgt_emb, tgt_mask=tgt_mask)
        return self.out(out)

# --- Training ---
def train(words_file):
    words = load_words(words_file)
    char2idx, idx2char = build_vocab(words)

    # +1 so that max index fits [0 .. len(char2idx)-1]
    vocab_size = len(char2idx)
    max_word = max(len(w) for w in words)
    max_len = max_word + 2  # <sos> + word + <eos>

    dataset = StringDataset(words, CURRENT_OBFUSCATOR, char2idx, max_len)
    loader  = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)

    model = TransformerModel(vocab_size, HIDDEN_SIZE, max_len).to(DEVICE)
    optim = torch.optim.Adam(model.parameters())
    criterion = nn.CrossEntropyLoss(ignore_index=0)

    os.makedirs(MODEL_DIR, exist_ok=True)

    for epoch in range(1, NUM_EPOCHS+1):
        model.train()
        total_loss = 0
        start = time.time()
        loop = tqdm(loader, desc=f"Epoch {epoch}/{NUM_EPOCHS}")
        for src, tgt in loop:
            src, tgt = src.to(DEVICE), tgt.to(DEVICE)
            optim.zero_grad()
            # teacher forcing: feed all but last token of target
            out = model(src, tgt[:, :-1])      # → (B, seq_len-1, vocab)
            logits = out.reshape(-1, vocab_size)
            gold   = tgt[:, 1:].reshape(-1)    # shift by 1
            loss   = criterion(logits, gold)
            loss.backward()
            optim.step()
            total_loss += loss.item()
            loop.set_postfix(loss=loss.item())

        print(f"Epoch {epoch} — avg loss {total_loss/len(loader):.4f} — time {time.time()-start:.1f}s")
        torch.save(model.state_dict(),  os.path.join(MODEL_DIR, f"epoch{epoch}.pth"))
        torch.save(char2idx,            os.path.join(MODEL_DIR, "char2idx.pth"))
        torch.save(idx2char,            os.path.join(MODEL_DIR, "idx2char.pth"))

    print("Training complete.")

# --- Main ---
if __name__ == "__main__":
    import sys
    if len(sys.argv) != 2:
        print("Usage: python train.py words.txt")
        exit(1)
    train(sys.argv[1])