{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":121150,"databundleVersionId":15833391}],"dockerImageVersionId":31329,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# ==========================================\n# BLOCK 1: IMPROVED TRAINING WITH ADAMW, WARMUP, INCREASED MAX_SEQ_LEN\n# ==========================================\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.nn.utils.rnn import pad_sequence\nfrom sklearn.model_selection import train_test_split\n\nimport pandas as pd\nimport numpy as np\nimport math\nimport re\nfrom collections import Counter\nfrom tqdm import tqdm\n\n# ==========================================\n# 1. CONFIGURATION (UPDATED)\n# ==========================================\nclass Config:\n    TRAIN_PATH = \"/kaggle/input/competitions/deep-past-initiative-machine-translation/train.csv\"\n    TEST_PATH = \"/kaggle/input/competitions/deep-past-initiative-machine-translation/test.csv\"\n    SUBMISSION_PATH = \"submission.csv\"\n    \n    MIN_WORD_FREQ = 1       \n    MAX_SEQ_LEN = 100       # INCREASED: 50 -> 100 for better handling of longer texts\n    VAL_SPLIT = 0.1\n    CHECKPOINT_PATH = \"transformer_best.pth\"\n    \n    PAD_IDX, SOS_IDX, EOS_IDX, UNK_IDX = 0, 1, 2, 3\n    PAD_TOKEN, SOS_TOKEN, EOS_TOKEN, UNK_TOKEN = \"<pad>\", \"<sos>\", \"<eos>\", \"<unk>\"\n    \n    EMBED_SIZE = 256        \n    NHEAD = 4               \n    FFN_HID_DIM = 512       \n    NUM_LAYERS = 3          \n    DROPOUT = 0.3        \n    \n    BATCH_SIZE = 128         \n    EPOCHS = 50            \n    LEARNING_RATE = 0.0005  \n    WEIGHT_DECAY = 0.01     # NEW: For AdamW optimizer\n    WARMUP_STEPS = 1000     # NEW: Learning rate warmup steps\n    \n    DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\ncfg = Config()\n\n# ==========================================\n# 2. TEXT PROCESSING & VOCABULARY\n# ==========================================\ndef clean_text(text, is_akkadian=True):\n    if pd.isna(text): return \"\"\n    text = str(text)\n    if is_akkadian:\n        text = re.sub(r'[!\\\\?\\\\/:\\\\.\\\\[\\\\]\\\\<\\\\>˹˺]', '', text)\n        text = text.replace('x', '<gap>').replace('...', '<big_gap>')\n    return text.lower().strip().split()\n\nclass Vocabulary:\n    def __init__(self, name):\n        self.word2idx = {cfg.PAD_TOKEN: cfg.PAD_IDX, cfg.SOS_TOKEN: cfg.SOS_IDX, \n                         cfg.EOS_TOKEN: cfg.EOS_IDX, cfg.UNK_TOKEN: cfg.UNK_IDX}\n        self.idx2word = {v: k for k, v in self.word2idx.items()}\n        self.word_counts = Counter()\n        \n    def add_sentence(self, tokens):\n        self.word_counts.update(tokens)\n            \n    def build_vocab(self):\n        idx = 4\n        for word, count in self.word_counts.items():\n            if count >= cfg.MIN_WORD_FREQ:\n                self.word2idx[word] = idx\n                self.idx2word[idx] = word\n                idx += 1\n                \n    def encode(self, tokens):\n        return [self.word2idx.get(word, cfg.UNK_IDX) for word in tokens]\n        \n    def __len__(self): return len(self.word2idx)\n\ndf = pd.read_csv(cfg.TRAIN_PATH)\ntrain_df, val_df = train_test_split(df, test_size=cfg.VAL_SPLIT, random_state=42)\n\nakk_vocab, eng_vocab = Vocabulary(\"akkadian\"), Vocabulary(\"english\")\n\nfor _, row in train_df.iterrows():\n    akk_vocab.add_sentence(clean_text(row['transliteration'], is_akkadian=True))\n    eng_vocab.add_sentence(clean_text(row['translation'], is_akkadian=False))\n\nakk_vocab.build_vocab()\neng_vocab.build_vocab()\n\n# ==========================================\n# 3. DATASET & DATALOADER\n# ==========================================\nclass TranslationDataset(Dataset):\n    def __init__(self, df, src_vocab, trg_vocab):\n        self.src_seqs, self.trg_seqs = [], []\n        for _, row in df.iterrows():\n            src_encoded = [cfg.SOS_IDX] + src_vocab.encode(clean_text(row['transliteration'], True)[:cfg.MAX_SEQ_LEN]) + [cfg.EOS_IDX]\n            trg_encoded = [cfg.SOS_IDX] + trg_vocab.encode(clean_text(row['translation'], False)[:cfg.MAX_SEQ_LEN]) + [cfg.EOS_IDX]\n            self.src_seqs.append(torch.tensor(src_encoded))\n            self.trg_seqs.append(torch.tensor(trg_encoded))\n            \n    def __len__(self): return len(self.src_seqs)\n    def __getitem__(self, idx): return self.src_seqs[idx], self.trg_seqs[idx]\n\ndef collate_fn(batch):\n    src_batch, trg_batch = zip(*batch)\n    return (pad_sequence(src_batch, padding_value=cfg.PAD_IDX, batch_first=True),\n            pad_sequence(trg_batch, padding_value=cfg.PAD_IDX, batch_first=True))\n\ntrain_loader = DataLoader(TranslationDataset(train_df, akk_vocab, eng_vocab), batch_size=cfg.BATCH_SIZE, shuffle=True, collate_fn=collate_fn)\nval_loader = DataLoader(TranslationDataset(val_df, akk_vocab, eng_vocab), batch_size=cfg.BATCH_SIZE, shuffle=False, collate_fn=collate_fn)\n\n# ==========================================\n# 4. TRANSFORMER ARCHITECTURE\n# ==========================================\ndef generate_square_subsequent_mask(sz, device):\n    mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(0, 1)\n    mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))\n    return mask\n\ndef create_mask(src, trg, pad_idx, device):\n    src_seq_len, trg_seq_len = src.shape[1], trg.shape[1]\n    trg_mask = generate_square_subsequent_mask(trg_seq_len, device)\n    src_mask = torch.zeros((src_seq_len, src_seq_len), device=device).type(torch.bool)\n    src_padding_mask = (src == pad_idx)\n    trg_padding_mask = (trg == pad_idx)\n    return src_mask, trg_mask, src_padding_mask, trg_padding_mask\n\nclass PositionalEncoding(nn.Module):\n    def __init__(self, emb_size, dropout, maxlen=5000):\n        super(PositionalEncoding, self).__init__()\n        den = torch.exp(-torch.arange(0, emb_size, 2) * math.log(10000) / emb_size)\n        pos = torch.arange(0, maxlen).reshape(maxlen, 1)\n        pos_embedding = torch.zeros((maxlen, emb_size))\n        pos_embedding[:, 0::2] = torch.sin(pos * den)\n        pos_embedding[:, 1::2] = torch.cos(pos * den)\n        self.dropout = nn.Dropout(dropout)\n        self.register_buffer('pos_embedding', pos_embedding.unsqueeze(0))\n\n    def forward(self, token_embedding):\n        return self.dropout(token_embedding + self.pos_embedding[:, :token_embedding.size(1), :])\n\nclass Seq2SeqTransformer(nn.Module):\n    def __init__(self, num_encoder_layers, num_decoder_layers, emb_size, nhead, src_vocab_size, trg_vocab_size, dim_feedforward, dropout):\n        super(Seq2SeqTransformer, self).__init__()\n        self.transformer = nn.Transformer(\n            d_model=emb_size, nhead=nhead,\n            num_encoder_layers=num_encoder_layers, num_decoder_layers=num_decoder_layers,\n            dim_feedforward=dim_feedforward, dropout=dropout, batch_first=True\n        )\n        self.src_tok_emb = nn.Embedding(src_vocab_size, emb_size, padding_idx=cfg.PAD_IDX)\n        self.trg_tok_emb = nn.Embedding(trg_vocab_size, emb_size, padding_idx=cfg.PAD_IDX)\n        self.positional_encoding = PositionalEncoding(emb_size, dropout=dropout)\n        self.generator = nn.Linear(emb_size, trg_vocab_size)\n        self.emb_size = emb_size\n\n    def forward(self, src, trg):\n        src_mask, trg_mask, src_padding_mask, trg_padding_mask = create_mask(src, trg, cfg.PAD_IDX, src.device)\n        \n        src_emb = self.positional_encoding(self.src_tok_emb(src) * math.sqrt(self.emb_size))\n        trg_emb = self.positional_encoding(self.trg_tok_emb(trg) * math.sqrt(self.emb_size))\n        \n        outs = self.transformer(src_emb, trg_emb, src_mask, trg_mask, None,\n                                src_padding_mask, trg_padding_mask, src_padding_mask)\n        return self.generator(outs)\n\n    def encode(self, src, src_mask):\n        return self.transformer.encoder(self.positional_encoding(self.src_tok_emb(src) * math.sqrt(self.emb_size)), src_mask)\n\n    def decode(self, trg, memory, trg_mask):\n        return self.transformer.decoder(self.positional_encoding(self.trg_tok_emb(trg) * math.sqrt(self.emb_size)), memory, trg_mask)\n\nmodel = Seq2SeqTransformer(cfg.NUM_LAYERS, cfg.NUM_LAYERS, cfg.EMBED_SIZE, cfg.NHEAD, \n                           len(akk_vocab), len(eng_vocab), cfg.FFN_HID_DIM, cfg.DROPOUT).to(cfg.DEVICE)\n\nfor p in model.parameters():\n    if p.dim() > 1: nn.init.xavier_uniform_(p)\n\nif torch.cuda.device_count() > 1:\n    model = nn.DataParallel(model)\n\n# ==========================================\n# 5. IMPROVED OPTIMIZER WITH ADAMW AND WARMUP\n# ==========================================\ndef get_optimizer_and_scheduler(model, cfg, total_steps):\n    # AdamW instead of Adam (better weight decay handling)\n    optimizer = optim.AdamW(\n        model.parameters(), \n        lr=cfg.LEARNING_RATE, \n        betas=(0.9, 0.98), \n        eps=1e-9,\n        weight_decay=cfg.WEIGHT_DECAY\n    )\n    \n    # Linear warmup + cosine annealing scheduler\n    def lr_lambda(step):\n        if step < cfg.WARMUP_STEPS:\n            # Linear warmup\n            return float(step) / float(max(1, cfg.WARMUP_STEPS))\n        else:\n            # Cosine annealing after warmup\n            progress = float(step - cfg.WARMUP_STEPS) / float(max(1, total_steps - cfg.WARMUP_STEPS))\n            return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress)))\n    \n    scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)\n    return optimizer, scheduler\n\n# ==========================================\n# 6. METRICS & TRAINING LOOP\n# ==========================================\ndef get_ngrams(sentence, n, is_char=False):\n    chars = sentence.replace(\" \", \"\") if is_char else sentence.split()\n    return Counter([\" \".join(chars[i:i+n]) if not is_char else chars[i:i+n] for i in range(len(chars)-n+1)])\n\ndef calculate_metrics(preds, refs):\n    matches_by_n, possible_by_n, total_pred_len, total_ref_len = [0]*4, [0]*4, 0, 0\n    c_match, c_pred, c_ref = [0]*6, [0]*6, [0]*6\n    w_match, w_pred, w_ref = [0]*2, [0]*2, [0]*2\n\n    for p, r in zip(preds, refs):\n        p_words, r_words = p.split(), r.split()\n        total_pred_len += len(p_words); total_ref_len += len(r_words)\n        \n        for n in range(1, 5):\n            p_ng, r_ng = get_ngrams(p, n), get_ngrams(r, n)\n            matches_by_n[n-1] += sum(min(count, r_ng[ng]) for ng, count in p_ng.items())\n            possible_by_n[n-1] += max(0, len(p_words) - n + 1)\n            \n        for n in range(1, 7):\n            p_c, r_c = get_ngrams(p, n, True), get_ngrams(r, n, True)\n            c_match[n-1] += sum(min(count, r_c[ng]) for ng, count in p_c.items())\n            c_pred[n-1] += sum(p_c.values()); c_ref[n-1] += sum(r_c.values())\n        for n in range(1, 3):\n            p_w, r_w = get_ngrams(p, n), get_ngrams(r, n)\n            w_match[n-1] += sum(min(count, r_w[ng]) for ng, count in p_w.items())\n            w_pred[n-1] += sum(p_w.values()); w_ref[n-1] += sum(r_w.values())\n\n    p_n = [max(m / p, 1e-10) if p > 0 else 1e-10 for m, p in zip(matches_by_n, possible_by_n)]\n    bp = 1.0 if total_pred_len > total_ref_len else (0.0 if total_pred_len == 0 else math.exp(1 - float(total_ref_len)/total_pred_len))\n    bleu = bp * math.exp(sum(math.log(p) for p in p_n) / 4) * 100\n\n    total_f, valid_ngrams = 0.0, 0\n    for match, pred, ref in zip(c_match + w_match, c_pred + w_pred, c_ref + w_ref):\n        if pred == 0 or ref == 0: continue\n        pr, re = match / pred, match / ref\n        if pr + re > 0: total_f += (5) * (pr * re) / ((4 * pr) + re) \n        valid_ngrams += 1\n    chrf = (total_f / valid_ngrams) * 100 if valid_ngrams > 0 else 0.0\n    return bleu, chrf, math.sqrt((bleu / 100) * (chrf / 100))\n\ndef tensor_to_text(tensor, vocab):\n    words = [vocab.idx2word[idx.item()] for idx in tensor if idx.item() not in [cfg.PAD_IDX, cfg.SOS_IDX, cfg.EOS_IDX, cfg.UNK_IDX]]\n    return \" \".join(words)\n\ncriterion = nn.CrossEntropyLoss(ignore_index=cfg.PAD_IDX)\ntotal_training_steps = len(train_loader) * cfg.EPOCHS\noptimizer, scheduler = get_optimizer_and_scheduler(model, cfg, total_training_steps)\n\ndef train_epoch_transformer(model, dataloader, optimizer, criterion, scheduler):\n    model.train()\n    epoch_loss = 0\n    for src, trg in tqdm(dataloader, desc=\"Training\", leave=False):\n        src, trg = src.to(cfg.DEVICE), trg.to(cfg.DEVICE)\n        \n        trg_input = trg[:, :-1]\n        trg_target = trg[:, 1:]\n        \n        optimizer.zero_grad()\n        logits = model(src, trg_input)\n        \n        loss = criterion(logits.reshape(-1, logits.shape[-1]), trg_target.reshape(-1))\n        loss.backward()\n        \n        # Gradient clipping for stability\n        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n        \n        optimizer.step()\n        scheduler.step()  # Update learning rate each step\n        epoch_loss += loss.item()\n    return epoch_loss / len(dataloader)\n\ndef evaluate_model(model, dataloader, eng_vocab):\n    model.eval()\n    preds_text, refs_text = [], []\n    base_model = model.module if isinstance(model, nn.DataParallel) else model\n    \n    with torch.no_grad():\n        for src, trg in tqdm(dataloader, desc=\"Validating\", leave=False):\n            src, trg = src.to(cfg.DEVICE), trg.to(cfg.DEVICE)\n            \n            src_mask = (torch.zeros(src.shape[1], src.shape[1])).type(torch.bool).to(cfg.DEVICE)\n            memory = base_model.encode(src, src_mask)\n            \n            ys = torch.ones(src.shape[0], 1).fill_(cfg.SOS_IDX).type(torch.long).to(cfg.DEVICE)\n            \n            for i in range(cfg.MAX_SEQ_LEN - 1):\n                trg_mask = generate_square_subsequent_mask(ys.size(1), cfg.DEVICE).type(torch.bool).to(cfg.DEVICE)\n                out = base_model.decode(ys, memory, trg_mask)\n                \n                prob = base_model.generator(out[:, -1, :])\n                prob[:, cfg.UNK_IDX] = float('-inf') \n                _, next_word = torch.max(prob, dim=1)\n                \n                ys = torch.cat([ys, next_word.unsqueeze(1)], dim=1)\n                \n            for i in range(ys.shape[0]):\n                preds_text.append(tensor_to_text(ys[i], eng_vocab))\n                refs_text.append(tensor_to_text(trg[i], eng_vocab))\n                \n    return calculate_metrics(preds_text, refs_text)\n\n# ==========================================\n# 7. TRAINING LOOP WITH IMPROVED CONFIG\n# ==========================================\nbest_score = 0.0\nprint(\"Starting Improved Training (AdamW + Warmup + MAX_SEQ_LEN=100)...\")\nprint(f\"Device: {cfg.DEVICE}\")\nprint(f\"Batch Size: {cfg.BATCH_SIZE}, Max Seq Len: {cfg.MAX_SEQ_LEN}\")\nprint(f\"Optimizer: AdamW (LR={cfg.LEARNING_RATE}, WD={cfg.WEIGHT_DECAY})\")\nprint(f\"Warmup Steps: {cfg.WARMUP_STEPS}, Total Steps: {total_training_steps}\")\n\nfor epoch in range(cfg.EPOCHS):\n    train_loss = train_epoch_transformer(model, train_loader, optimizer, criterion, scheduler)\n    bleu, chrf, score = evaluate_model(model, val_loader, eng_vocab)\n    \n    current_lr = optimizer.param_groups[0]['lr']\n    print(f\"Epoch {epoch+1:03d}/{cfg.EPOCHS} | LR: {current_lr:.6f} | Loss: {train_loss:.4f} | BLEU: {bleu:.2f} | chrF++: {chrf:.2f} | Score: {score:.4f}\")\n    \n    if score > best_score:\n        best_score = score\n        print(f\"--> New best score! Saving to {cfg.CHECKPOINT_PATH}\")\n        torch.save(model.module.state_dict() if isinstance(model, nn.DataParallel) else model.state_dict(), cfg.CHECKPOINT_PATH)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ==========================================\n# BLOCK 2: INFERENCE & SUBMISSION (UPDATED FOR IMPROVED MODEL)\n# ==========================================\n\n# ==========================================\n# 6. INFERENCE & SUBMISSION\n# ==========================================\ndef translate_sentence_transformer(model, sentence_text, src_vocab, trg_vocab, device):\n    model.eval()\n    base_model = model.module if isinstance(model, nn.DataParallel) else model\n    \n    tokens = clean_text(sentence_text, is_akkadian=True)\n    src_indices = [src_vocab.word2idx.get(token, cfg.UNK_IDX) for token in tokens]\n    \n    # Uses updated MAX_SEQ_LEN=100\n    src_tensor = torch.LongTensor([cfg.SOS_IDX] + src_indices + [cfg.EOS_IDX]).unsqueeze(0).to(device)\n    src_mask = (torch.zeros(src_tensor.shape[1], src_tensor.shape[1])).type(torch.bool).to(device)\n    \n    with torch.no_grad():\n        memory = base_model.encode(src_tensor, src_mask)\n    \n    ys = torch.ones(1, 1).fill_(cfg.SOS_IDX).type(torch.long).to(device)\n    \n    for _ in range(cfg.MAX_SEQ_LEN):\n        trg_mask = generate_square_subsequent_mask(ys.size(1), device)\n        \n        with torch.no_grad():\n            out = base_model.decode(ys, memory, trg_mask)\n            prob = base_model.generator(out[:, -1, :])\n            prob[0, cfg.UNK_IDX] = float('-inf')\n            \n            next_word = prob.argmax(1).item()\n            \n        ys = torch.cat([ys, torch.ones(1, 1).fill_(next_word).type(torch.long).to(device)], dim=1)\n        if next_word == cfg.EOS_IDX: break\n            \n    words = [trg_vocab.idx2word[idx.item()] for idx in ys[0] if idx.item() not in [cfg.SOS_IDX, cfg.EOS_IDX, cfg.PAD_IDX]]\n    final_text = \" \".join(words)\n    return final_text[0].upper() + final_text[1:] + \".\" if final_text else \"\"\n\nprint(\"\\nLoading best model weights for Inference...\")\nif isinstance(model, nn.DataParallel):\n    model.module.load_state_dict(torch.load(cfg.CHECKPOINT_PATH))\nelse:\n    model.load_state_dict(torch.load(cfg.CHECKPOINT_PATH))\n\ntest_df = pd.read_csv(cfg.TEST_PATH, dtype={'line_start': str, 'line_end': str})\npredictions = []\n\nfor _, row in tqdm(test_df.iterrows(), total=len(test_df), desc=\"Translating\"):\n    predictions.append(translate_sentence_transformer(model, str(row['transliteration']), akk_vocab, eng_vocab, cfg.DEVICE))\n\n# Create submission file format required by Kaggle\nsubmission = pd.DataFrame({\n    'id': test_df['id'],\n    'translation': predictions\n})\n\n# Sanitize empty predictions\nsubmission['translation'] = submission['translation'].replace('', 'Empty translation.')\nsubmission['translation'] = submission['translation'].fillna('Empty translation.')\nsubmission['translation'] = submission['translation'].astype(str)\n\nsubmission.to_csv(cfg.SUBMISSION_PATH, index=False)\nprint(f\"\\n✅ Submission saved to {cfg.SUBMISSION_PATH}\")\nprint(f\"📊 Generated {len(predictions)} translations using improved model (MAX_SEQ_LEN={cfg.MAX_SEQ_LEN})\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}