Add TabuLM training and evaluation code

code/adamw.py

@@ -0,0 +1,109 @@
+import math
+
+import torch
+from torch.optim import Optimizer
+
+# From: https://github.com/uds-lsv/bert-stable-fine-tuning
+# Paper: On the Stability of Fine-tuning BERT: Misconceptions, Explanations, and Strong Baselines
+# Marius Mosbach, Maksym Andriushchenko, Dietrich Klakow
+# https://arxiv.org/abs/2006.04884
+
+class mAdamW(Optimizer):
+    """ Implements Adam algorithm with weight decay fix.
+    Parameters:
+        lr (float): learning rate. Default 1e-3.
+        betas (tuple of 2 floats): Adams beta parameters (b1, b2). Default: (0.9, 0.999)
+        eps (float): Adams epsilon. Default: 1e-6
+        weight_decay (float): Weight decay. Default: 0.0
+        correct_bias (bool): can be set to False to avoid correcting bias in Adam (e.g. like in Bert TF repository). Default True.
+        local_normalization (bool):
+        max_grad_norm (bool):
+    """
+
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-6, weight_decay=0.0, correct_bias=True,
+                 local_normalization=False, max_grad_norm=-1):
+        if lr < 0.0:
+            raise ValueError("Invalid learning rate: {} - should be >= 0.0".format(lr))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError("Invalid beta parameter: {} - should be in [0.0, 1.0[".format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError("Invalid beta parameter: {} - should be in [0.0, 1.0[".format(betas[1]))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {} - should be >= 0.0".format(eps))
+        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, correct_bias=correct_bias,
+                        local_normalization=local_normalization, max_grad_norm=max_grad_norm)
+        super().__init__(params, defaults)
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        updates = []
+
+        for group in self.param_groups:
+
+            group_updates = []
+
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+
+                # Normalize gradients locally (layer-wise)
+                if group["local_normalization"]:
+                    torch.nn.utils.clip_grad_norm_(p, group["max_grad_norm"])
+
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError("Adam does not support sparse gradients, please consider SparseAdam instead")
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state["step"] = 0
+                    # Exponential moving average of gradient values
+                    state["exp_avg"] = torch.zeros_like(p.data)
+                    # Exponential moving average of squared gradient values
+                    state["exp_avg_sq"] = torch.zeros_like(p.data)
+
+                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+                beta1, beta2 = group["betas"]
+
+                state["step"] += 1
+
+                # Decay the first and second moment running average coefficient
+                # In-place operations to update the averages at the same time
+                exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1)
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
+                denom = exp_avg_sq.sqrt().add_(group["eps"])
+
+                step_size = group["lr"]
+                if group["correct_bias"]:  # No bias correction for Bert
+                    bias_correction1 = 1.0 - beta1 ** state["step"]
+                    bias_correction2 = 1.0 - beta2 ** state["step"]
+                    step_size = step_size * math.sqrt(bias_correction2) / bias_correction1
+
+                p.data.addcdiv_(-step_size, exp_avg, denom)
+
+                # Just adding the square of the weights to the loss function is *not*
+                # the correct way of using L2 regularization/weight decay with Adam,
+                # since that will interact with the m and v parameters in strange ways.
+                #
+                # Instead we want to decay the weights in a manner that doesn't interact
+                # with the m/v parameters. This is equivalent to adding the square
+                # of the weights to the loss with plain (non-momentum) SGD.
+                # Add weight decay at the end (fixed version)
+                if group["weight_decay"] > 0.0:
+                    p.data.add_(p.data, alpha=-group["lr"] * group["weight_decay"])
+
+                group_updates.append((exp_avg, denom))
+
+            updates.append(group_updates)
+
+        return loss, updates

code/eval_baselines.py

@@ -0,0 +1,423 @@
+#!/usr/bin/env python3
+"""Baseline evaluation on TabQA-kin using mBERT, XLM-R, and KinyaBERT.
+
+Each baseline uses the same cell-selection approach as TabuLM:
+- Serialize table as flat text, track (row, col) token positions
+- Pool hidden states per cell → linear cell scorer
+- Same 80/20 train/dev split (seed=42), 20 epochs, lr=2e-5
+- Gold cell found by text-matching (same as finetune_tabqa.py)
+"""
+
+import argparse
+import json
+import os
+import random
+import sys
+from datetime import datetime
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn.utils import clip_grad_norm_
+from transformers import AutoTokenizer, AutoModel
+
+DATA_DIR = '/shared/scratch/0/tmp/v_ireddi_rakshitha_results/tabulm/data'
+CODE_DIR = '/shared/scratch/0/tmp/v_ireddi_rakshitha_results/tabulm/code'
+TABQA_FILE = os.path.join(DATA_DIR, 'tabqa_kin.json')
+CSV_DIR    = os.path.join(DATA_DIR, 'tables')
+
+sys.path.insert(0, CODE_DIR)
+from tabular_serializer import serialize_csv, TableCell
+
+# ── Model registry ─────────────────────────────────────────────────────────────
+MODELS = {
+    'mbert':     'bert-base-multilingual-cased',
+    'xlmr':      'xlm-roberta-base',
+    'kinyabert': 'jean-paul/KinyaBERT-large',
+}
+
+
+# ── Gold cell lookup (same as finetune_tabqa.py) ───────────────────────────────
+
+def find_gold_cell(cells: List[TableCell], answer_text: str,
+                   question_text: str = '') -> Optional[Tuple[int, int]]:
+    answer_norm = answer_text.strip().lower()
+    matches = [(c.row_id, c.col_id) for c in cells
+               if c.row_id > 1 and c.col_id > 0 and c.content.strip() == answer_text.strip()]
+    if not matches:
+        matches = [(c.row_id, c.col_id) for c in cells
+                   if c.row_id > 1 and c.col_id > 0 and c.content.strip().lower() == answer_norm]
+    if not matches:
+        return None
+    if len(matches) == 1:
+        return matches[0]
+    if question_text:
+        q_words = set(question_text.lower().split())
+        row_labels  = {c.row_id: c.content.strip().lower()
+                       for c in cells if c.col_id == 1 and c.row_id > 1}
+        col_headers = {c.col_id: c.content.strip().lower()
+                       for c in cells if c.row_id == 1 and c.col_id > 0}
+        best, best_score = None, (-1, -1)
+        for (row_id, col_id) in matches:
+            r_sc = len(q_words & set(row_labels.get(row_id,  '').split()))
+            c_sc = len(q_words & set(col_headers.get(col_id, '').split()))
+            if (r_sc, c_sc) > best_score:
+                best_score, best = (r_sc, c_sc), (row_id, col_id)
+        if best_score[0] > 0 or best_score[1] > 0:
+            return best
+    return min(matches, key=lambda rc: (rc[0], rc[1]))
+
+
+# ── Table serialization for HF tokenizers ─────────────────────────────────────
+
+def serialize_table_for_bert(cells: List[TableCell], question: str,
+                              tokenizer, max_seq_len: int = 512):
+    """
+    Linearize table + question for a BERT-style model.
+    Format: [CLS] h1 | h2 | ... [SEP] v1 | v2 | ... [SEP] ... question [SEP]
+    Tracks which token positions belong to each (row_id, col_id) data cell.
+    Returns (input_ids, attn_mask, ordered_cells, cell_to_positions) or None.
+    """
+    rows: Dict[int, Dict[int, str]] = {}
+    for c in cells:
+        if c.row_id > 0 and c.col_id > 0:
+            rows.setdefault(c.row_id, {})[c.col_id] = c.content.strip()
+
+    if not rows:
+        return None
+
+    sorted_rows = sorted(rows.keys())
+    col_ids_sorted = sorted(rows[sorted_rows[0]].keys())
+
+    tokens: List[int] = [tokenizer.cls_token_id]
+    cell_to_positions: Dict[Tuple[int, int], List[int]] = {}
+
+    pipe_id = tokenizer.convert_tokens_to_ids('|')
+    if pipe_id == tokenizer.unk_token_id:
+        pipe_id = None   # skip pipe separator if not in vocab
+
+    for row_id in sorted_rows:
+        for i, col_id in enumerate(col_ids_sorted):
+            content = rows[row_id].get(col_id, '')
+            sub = tokenizer.encode(content, add_special_tokens=False)
+            if sub:
+                start = len(tokens)
+                tokens.extend(sub)
+                if row_id > 1:   # data rows only (header row_id=1 not scored)
+                    cell_to_positions.setdefault((row_id, col_id), []).extend(
+                        range(start, len(tokens))
+                    )
+            if pipe_id is not None and i < len(col_ids_sorted) - 1:
+                tokens.append(pipe_id)
+        tokens.append(tokenizer.sep_token_id)
+
+    # Question
+    q_ids = tokenizer.encode(question, add_special_tokens=False)
+    tokens.extend(q_ids)
+    tokens.append(tokenizer.sep_token_id)
+
+    # Truncate
+    if len(tokens) > max_seq_len:
+        tokens = tokens[:max_seq_len - 1] + [tokenizer.sep_token_id]
+        cell_to_positions = {
+            k: [p for p in v if p < max_seq_len]
+            for k, v in cell_to_positions.items()
+        }
+        cell_to_positions = {k: v for k, v in cell_to_positions.items() if v}
+
+    input_ids = torch.tensor(tokens, dtype=torch.long)
+    attn_mask = torch.ones(len(tokens), dtype=torch.long)
+    ordered_cells = sorted(cell_to_positions.keys(),
+                           key=lambda rc: min(cell_to_positions[rc]))
+
+    return input_ids, attn_mask, ordered_cells, cell_to_positions
+
+
+# ── Cell-scoring model ─────────────────────────────────────────────────────────
+
+class BaselineCellScorer(nn.Module):
+    def __init__(self, encoder, d_model: int):
+        super().__init__()
+        self.encoder   = encoder
+        self.cell_head = nn.Linear(d_model, 1)
+        nn.init.normal_(self.cell_head.weight, std=0.02)
+        nn.init.zeros_(self.cell_head.bias)
+
+    def forward(self, input_ids, attention_mask, ordered_cells,
+                cell_to_positions, device):
+        out = self.encoder(
+            input_ids=input_ids.unsqueeze(0).to(device),
+            attention_mask=attention_mask.unsqueeze(0).to(device),
+        )
+        hidden = out.last_hidden_state[0]   # (S, d)
+        S = hidden.size(0)
+
+        cell_embeds, valid_cells = [], []
+        for rc in ordered_cells:
+            positions = [p for p in cell_to_positions[rc] if p < S]
+            if not positions:
+                continue
+            h = hidden[positions].mean(0)
+            cell_embeds.append(h)
+            valid_cells.append(rc)
+
+        if not cell_embeds:
+            return None, []
+
+        scores = self.cell_head(torch.stack(cell_embeds)).squeeze(-1)
+        return scores, valid_cells
+
+
+# ── Lookup prediction helper ───────────────────────────────────────────────────
+
+def _predict_lookup(scores, valid_cells, cells, question_text):
+    q_words = set(question_text.lower().split())
+    row_labels  = {c.row_id: c.content.strip().lower()
+                   for c in cells if c.col_id == 1 and c.row_id > 1}
+    col_headers = {c.col_id: c.content.strip().lower()
+                   for c in cells if c.row_id == 1 and c.col_id > 0}
+
+    if not q_words:
+        return valid_cells[scores.argmax().item()]
+
+    row_score = {r: len(q_words & set(lbl.split())) for r, lbl in row_labels.items()}
+    col_score = {c: len(q_words & set(hdr.split())) for c, hdr in col_headers.items()}
+
+    best_row = max(row_score, key=row_score.get) if row_score else None
+    best_col = max(col_score, key=col_score.get) if col_score else None
+    has_row  = best_row is not None and row_score[best_row] > 0
+    has_col  = best_col is not None and col_score[best_col] > 0
+
+    if not has_row and not has_col:
+        return valid_cells[scores.argmax().item()]
+
+    if has_row:
+        row_indices = [i for i, rc in enumerate(valid_cells) if rc[0] == best_row]
+    else:
+        row_indices = list(range(len(valid_cells)))
+
+    if not row_indices:
+        return valid_cells[scores.argmax().item()]
+
+    if has_col:
+        both = [i for i in row_indices if valid_cells[i][1] == best_col]
+        if both:
+            return valid_cells[max(both, key=lambda i: scores[i].item())]
+
+    return valid_cells[max(row_indices, key=lambda i: scores[i].item())]
+
+
+def _predict_comparison(scores, valid_cells, cells, question_text):
+    """For comparison questions: restrict to col-1 (entity name) cells of the
+    top-2 question-relevant rows. The answer is the winning entity name."""
+    q_words = set(question_text.lower().split())
+    row_labels = {c.row_id: c.content.strip().lower()
+                  for c in cells if c.col_id == 1 and c.row_id > 1}
+
+    if not q_words or not row_labels:
+        return valid_cells[scores.argmax().item()]
+
+    row_score = {r: len(q_words & set(lbl.split())) for r, lbl in row_labels.items()}
+    top2 = sorted([r for r, s in row_score.items() if s > 0],
+                  key=lambda r: row_score[r], reverse=True)[:2]
+
+    if not top2:
+        return valid_cells[scores.argmax().item()]
+
+    cand = [i for i, (r, c) in enumerate(valid_cells) if r in top2 and c == 1]
+    if not cand:
+        cand = [i for i, (r, c) in enumerate(valid_cells) if r in top2]
+    if not cand:
+        return valid_cells[scores.argmax().item()]
+
+    return valid_cells[max(cand, key=lambda i: scores[i].item())]
+
+
+# ── Evaluation ─────────────────────────────────────────────────────────────────
+
+def evaluate(model, tokenizer, items, csv_dir, device, split='dev'):
+    model.eval()
+    correct, total, skipped = 0, 0, 0
+    by_type: Dict[str, List[int]] = {}
+
+    with torch.no_grad():
+        for item in items:
+            csv_path = os.path.join(csv_dir, item['table_file'])
+            if not os.path.exists(csv_path):
+                skipped += 1; continue
+            cells = serialize_csv(csv_path)
+            if not cells:
+                skipped += 1; continue
+
+            question = item['question']
+            atype    = item.get('answer_type', '?')
+
+            gold_rc = find_gold_cell(cells, item['answer'], question_text=question)
+            if gold_rc is None:
+                skipped += 1; continue
+
+            enc = serialize_table_for_bert(cells, question, tokenizer)
+            if enc is None:
+                skipped += 1; continue
+            input_ids, attn_mask, ordered_cells, cell_to_positions = enc
+
+            if gold_rc not in cell_to_positions:
+                skipped += 1; continue
+
+            scores, valid_cells = model(input_ids, attn_mask, ordered_cells,
+                                        cell_to_positions, device)
+            if scores is None or gold_rc not in valid_cells:
+                skipped += 1; continue
+
+            if atype == 'lookup':
+                pred_rc = _predict_lookup(scores, valid_cells, cells, question)
+            elif atype == 'comparison':
+                pred_rc = _predict_comparison(scores, valid_cells, cells, question)
+            else:
+                pred_rc = valid_cells[scores.argmax().item()]
+            hit = int(pred_rc == gold_rc)
+            correct += hit; total += 1
+            by_type.setdefault(atype, []).append(hit)
+
+    em = correct / total if total > 0 else 0.0
+    ts = datetime.now().strftime('%H:%M:%S')
+    print(f'  [{ts}] {split} EM={em:.4f} ({correct}/{total}, {skipped} skipped)')
+    for t, hits in sorted(by_type.items()):
+        print(f'    {t}: {sum(hits)}/{len(hits)} = {sum(hits)/len(hits):.3f}')
+    return em, {t: round(sum(h)/len(h), 4) for t, h in by_type.items()}
+
+
+# ── Training ───────────────────────────────────────────────────────────────────
+
+def train_baseline(model_name, hf_name, train_items, dev_items, csv_dir,
+                   device, num_epochs=20, lr=2e-5):
+    print(f'\n[BL] ===== {model_name.upper()} ({hf_name}) =====')
+    tokenizer = AutoTokenizer.from_pretrained(hf_name)
+    encoder   = AutoModel.from_pretrained(hf_name)
+    d_model   = encoder.config.hidden_size
+    model     = BaselineCellScorer(encoder, d_model).to(device)
+
+    print(f'[BL] d_model={d_model}  '
+          f'params={sum(p.numel() for p in model.parameters()):,}')
+
+    print('[BL] Zero-shot baseline:')
+    best_em, best_by_type = evaluate(model, tokenizer, dev_items,
+                                     csv_dir, device, 'zero-shot')
+
+    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
+    criterion = nn.CrossEntropyLoss()
+    history   = []
+
+    for epoch in range(1, num_epochs + 1):
+        model.train()
+        random.shuffle(train_items)
+        total_loss, n_ok, n_skip = 0.0, 0, 0
+
+        for item in train_items:
+            csv_path = os.path.join(csv_dir, item['table_file'])
+            if not os.path.exists(csv_path):
+                n_skip += 1; continue
+            cells = serialize_csv(csv_path)
+            if not cells:
+                n_skip += 1; continue
+
+            gold_rc = find_gold_cell(cells, item['answer'])
+            if gold_rc is None:
+                n_skip += 1; continue
+
+            enc = serialize_table_for_bert(cells, item['question'], tokenizer)
+            if enc is None:
+                n_skip += 1; continue
+            input_ids, attn_mask, ordered_cells, cell_to_positions = enc
+
+            if gold_rc not in cell_to_positions:
+                n_skip += 1; continue
+
+            scores, valid_cells = model(input_ids, attn_mask, ordered_cells,
+                                        cell_to_positions, device)
+            if scores is None or gold_rc not in valid_cells:
+                n_skip += 1; continue
+
+            gold_idx = torch.tensor([valid_cells.index(gold_rc)],
+                                    dtype=torch.long, device=device)
+            loss = criterion(scores.unsqueeze(0), gold_idx)
+            optimizer.zero_grad()
+            loss.backward()
+            clip_grad_norm_(model.parameters(), 1.0)
+            optimizer.step()
+
+            total_loss += loss.item(); n_ok += 1
+
+        ts  = datetime.now().strftime('%H:%M:%S')
+        avg = total_loss / max(n_ok, 1)
+        print(f'[{ts}] Epoch {epoch}/{num_epochs}  loss={avg:.4f}  '
+              f'trained={n_ok}  skipped={n_skip}')
+
+        em, by_type = evaluate(model, tokenizer, dev_items, csv_dir, device)
+        history.append({'epoch': epoch, 'loss': avg, 'dev_em': em})
+
+        if em > best_em:
+            best_em, best_by_type = em, by_type
+            sp = os.path.join(DATA_DIR, f'baseline_{model_name}_best.pt')
+            torch.save({'model_state_dict': model.state_dict(), 'em': em}, sp)
+            print(f'  ** New best EM={best_em:.4f}')
+
+    print(f'\n[BL] {model_name.upper()} done. Best EM={best_em:.4f}')
+    return best_em, best_by_type, history
+
+
+# ── Main ───────────────────────────────────────────────────────────────────────
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--models', nargs='+', default=['mbert'],
+                        choices=list(MODELS.keys()))
+    parser.add_argument('--epochs', type=int, default=20)
+    parser.add_argument('--lr',     type=float, default=2e-5)
+    args = parser.parse_args()
+
+    random.seed(42); np.random.seed(42); torch.manual_seed(42)
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f'[BL] Device: {device}')
+
+    with open(TABQA_FILE) as f:
+        all_items = json.load(f)
+    random.shuffle(all_items)
+    split_n     = int(0.8 * len(all_items))
+    train_items = all_items[:split_n]
+    dev_items   = all_items[split_n:]
+    print(f'[BL] {len(all_items)} items  Train={len(train_items)}  Dev={len(dev_items)}')
+
+    all_results = {}
+    for model_name in args.models:
+        hf_name = MODELS[model_name]
+        best_em, best_by_type, history = train_baseline(
+            model_name, hf_name, train_items, dev_items, CSV_DIR, device,
+            num_epochs=args.epochs, lr=args.lr,
+        )
+        all_results[model_name] = {
+            'hf_name': hf_name, 'best_em': best_em,
+            'by_type': best_by_type, 'history': history,
+        }
+
+    out = os.path.join(DATA_DIR, 'baseline_results.json')
+    with open(out, 'w') as f:
+        json.dump(all_results, f, indent=2)
+    print(f'\n[BL] Results saved to {out}')
+
+    print('\n' + '='*58)
+    print(f'{"Model":<18} {"EM":>6}  {"Agg":>5}  {"Cmp":>5}  {"Lkp":>5}')
+    print('-'*58)
+    for name, r in all_results.items():
+        bt = r['by_type']
+        print(f'{name:<18} {r["best_em"]:>6.3f}  '
+              f'{bt.get("aggregation",0):>5.3f}  '
+              f'{bt.get("comparison",0):>5.3f}  '
+              f'{bt.get("lookup",0):>5.3f}')
+    print(f'{"TabuLM (ours)":<18} {"0.560":>6}  '
+          f'{"0.792":>5}  {"0.417":>5}  {"0.286":>5}')
+
+
+if __name__ == '__main__':
+    main()

code/eval_compare.py

@@ -0,0 +1,377 @@
+#!/usr/bin/env python3
+"""Paired per-item evaluation: TabuLM vs KinyaBERT on the same dev items.
+
+Outputs:
+  data/compare_results.json  — per-item predictions for both models
+  Prints bootstrap 95% CI for both models
+  Prints comparison examples (TabuLM correct, KinyaBERT wrong) for error analysis
+"""
+
+import argparse
+import json
+import os
+import random
+import sys
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+from transformers import AutoTokenizer, AutoModel
+
+DATA_DIR   = '/shared/scratch/0/tmp/v_ireddi_rakshitha_results/tabulm/data'
+CODE_DIR   = '/shared/scratch/0/tmp/v_ireddi_rakshitha_results/tabulm/code'
+TABQA_FILE = os.path.join(DATA_DIR, 'tabqa_kin.json')
+CSV_DIR    = os.path.join(DATA_DIR, 'tables')
+
+TABULM_CKPT  = os.path.join(DATA_DIR, 'finetune_tabqa_v3_best.pt')
+KINYABERT_CKPT = os.path.join(DATA_DIR, 'baseline_kinyabert_best.pt')
+
+sys.path.insert(0, CODE_DIR)
+
+import youtokentome as yttm
+from morpho_data_loaders import KBVocab
+from tabular_serializer import serialize_csv, table_cells_to_text, TableCell
+from morpho_stub import parse_text_stub
+from tabulm_model import TabuLM, tabulm_base
+
+
+# ── Shared gold-cell lookup ────────────────────────────────────────────────────
+
+def find_gold_cell(cells: List[TableCell], answer_text: str,
+                   question_text: str = '') -> Optional[Tuple[int, int]]:
+    answer_norm = answer_text.strip().lower()
+    matches = [(c.row_id, c.col_id) for c in cells
+               if c.row_id > 1 and c.col_id > 0
+               and c.content.strip() == answer_text.strip()]
+    if not matches:
+        matches = [(c.row_id, c.col_id) for c in cells
+                   if c.row_id > 1 and c.col_id > 0
+                   and c.content.strip().lower() == answer_norm]
+    if not matches:
+        return None
+    if len(matches) == 1:
+        return matches[0]
+    if question_text:
+        q_words = set(question_text.lower().split())
+        row_labels  = {c.row_id: c.content.strip().lower()
+                       for c in cells if c.col_id == 1 and c.row_id > 1}
+        col_headers = {c.col_id: c.content.strip().lower()
+                       for c in cells if c.row_id == 1 and c.col_id > 0}
+        best, best_score = None, (-1, -1)
+        for (row_id, col_id) in matches:
+            r_sc = len(q_words & set(row_labels.get(row_id,  '').split()))
+            c_sc = len(q_words & set(col_headers.get(col_id, '').split()))
+            if (r_sc, c_sc) > best_score:
+                best_score, best = (r_sc, c_sc), (row_id, col_id)
+        if best_score[0] > 0 or best_score[1] > 0:
+            return best
+    return min(matches, key=lambda rc: (rc[0], rc[1]))
+
+
+def _predict_lookup(scores, valid_cells, cells, question_text):
+    q_words = set(question_text.lower().split())
+    row_labels  = {c.row_id: c.content.strip().lower()
+                   for c in cells if c.col_id == 1 and c.row_id > 1}
+    col_headers = {c.col_id: c.content.strip().lower()
+                   for c in cells if c.row_id == 1 and c.col_id > 0}
+    if not q_words:
+        return valid_cells[scores.argmax().item()]
+    row_score = {r: len(q_words & set(lbl.split())) for r, lbl in row_labels.items()}
+    col_score = {c: len(q_words & set(hdr.split())) for c, hdr in col_headers.items()}
+    best_row = max(row_score, key=row_score.get) if row_score else None
+    best_col = max(col_score, key=col_score.get) if col_score else None
+    has_row  = best_row is not None and row_score[best_row] > 0
+    has_col  = best_col is not None and col_score[best_col] > 0
+    if not has_row and not has_col:
+        return valid_cells[scores.argmax().item()]
+    row_indices = [i for i, rc in enumerate(valid_cells) if rc[0] == best_row] if has_row \
+                  else list(range(len(valid_cells)))
+    if not row_indices:
+        return valid_cells[scores.argmax().item()]
+    if has_col:
+        both = [i for i in row_indices if valid_cells[i][1] == best_col]
+        if both:
+            return valid_cells[max(both, key=lambda i: scores[i].item())]
+    return valid_cells[max(row_indices, key=lambda i: scores[i].item())]
+
+
+# ── TabuLM inference ───────────────────────────────────────────────────────────
+
+def load_tabulm(device):
+    from finetune_tabqa import (detect_arch_from_state, encode_item, TabQAModel)
+    bpe = yttm.BPE(model=os.path.join(DATA_DIR, 'BPE-30k.mdl'))
+    kv  = KBVocab()
+    kv.load_state_dict(torch.load(
+        os.path.join(DATA_DIR, 'kb_vocab_state_dict_2021-02-07.pt'),
+        map_location='cpu'))
+
+    saved = torch.load(TABULM_CKPT, map_location='cpu')
+    state = saved['model_state_dict']
+    pretrain_state = {k: v for k, v in state.items()
+                      if not k.startswith('cell_head')}
+    if all(k.startswith('module.') for k in pretrain_state):
+        pretrain_state = {k[7:]: v for k, v in pretrain_state.items()}
+
+    args = detect_arch_from_state(pretrain_state)
+    tabulm_obj = tabulm_base(kv, None, None, device, args, saved_model_file=None)
+    model = TabQAModel(tabulm_obj).to(device)
+    model.load_state_dict(saved['model_state_dict'], strict=True)
+    model.eval()
+    return model, args, kv, bpe, encode_item
+
+
+def tabulm_predict(model, args, kv, bpe, encode_item, cells, question, atype, device):
+    from finetune_tabqa import encode_item as _enc
+    enc = _enc(cells, question, kv, bpe)
+    if enc is None:
+        return None
+    (pos_tags, stems, affixes, tokens_lengths,
+     row_ids, col_ids, cell_types, ordered_cells, cell_to_positions) = enc
+    with torch.no_grad():
+        scores, valid_cells = model(
+            args, pos_tags, stems, affixes,
+            tokens_lengths, row_ids, col_ids, cell_types,
+            ordered_cells, cell_to_positions, device)
+    if scores is None:
+        return None
+    if atype == 'lookup':
+        return _predict_lookup(scores, valid_cells, cells, question)
+    return valid_cells[scores.argmax().item()]
+
+
+# ── KinyaBERT inference ────────────────────────────────────────────────────────
+
+class KBCellScorer(nn.Module):
+    def __init__(self, encoder, d_model):
+        super().__init__()
+        self.encoder   = encoder
+        self.cell_head = nn.Linear(d_model, 1)
+
+    def forward(self, input_ids, attention_mask, ordered_cells,
+                cell_to_positions, device):
+        out    = self.encoder(
+            input_ids=input_ids.unsqueeze(0).to(device),
+            attention_mask=attention_mask.unsqueeze(0).to(device))
+        hidden = out.last_hidden_state[0]
+        S      = hidden.size(0)
+        cell_embeds, valid_cells = [], []
+        for rc in ordered_cells:
+            positions = [p for p in cell_to_positions[rc] if p < S]
+            if not positions:
+                continue
+            cell_embeds.append(hidden[positions].mean(0))
+            valid_cells.append(rc)
+        if not cell_embeds:
+            return None, []
+        scores = self.cell_head(torch.stack(cell_embeds)).squeeze(-1)
+        return scores, valid_cells
+
+
+def load_kinyabert(device):
+    from eval_baselines import serialize_table_for_bert
+    hf_name   = 'jean-paul/KinyaBERT-large'
+    tokenizer = AutoTokenizer.from_pretrained(hf_name)
+    encoder   = AutoModel.from_pretrained(hf_name)
+    d_model   = encoder.config.hidden_size
+    model     = KBCellScorer(encoder, d_model).to(device)
+    saved     = torch.load(KINYABERT_CKPT, map_location='cpu')
+    model.load_state_dict(saved['model_state_dict'], strict=False)
+    model.eval()
+    return model, tokenizer, serialize_table_for_bert
+
+
+def kb_predict(model, tokenizer, serialize_fn, cells, question, atype, device):
+    enc = serialize_fn(cells, question, tokenizer)
+    if enc is None:
+        return None
+    input_ids, attn_mask, ordered_cells, cell_to_positions = enc
+    with torch.no_grad():
+        scores, valid_cells = model(
+            input_ids, attn_mask, ordered_cells, cell_to_positions, device)
+    if scores is None:
+        return None
+    if atype == 'lookup':
+        return _predict_lookup(scores, valid_cells, cells, question)
+    return valid_cells[scores.argmax().item()]
+
+
+# ── Bootstrap CI ──────────────────────────────────────────────────────────────
+
+def bootstrap_ci(hits: List[int], n_boot: int = 10000, alpha: float = 0.05):
+    hits_arr = np.array(hits, dtype=float)
+    n = len(hits_arr)
+    boot_means = [np.mean(np.random.choice(hits_arr, size=n, replace=True))
+                  for _ in range(n_boot)]
+    boot_means = np.array(boot_means)
+    lo = np.percentile(boot_means, 100 * alpha / 2)
+    hi = np.percentile(boot_means, 100 * (1 - alpha / 2))
+    return float(np.mean(hits_arr)), lo, hi
+
+
+def paired_bootstrap_pvalue(hits_a: List[int], hits_b: List[int],
+                             n_boot: int = 10000) -> float:
+    """One-sided test: P(A > B) under H0 that A and B have same mean."""
+    a = np.array(hits_a, dtype=float)
+    b = np.array(hits_b, dtype=float)
+    n = len(a)
+    assert len(b) == n
+    obs_diff = np.mean(a) - np.mean(b)
+    diffs = []
+    for _ in range(n_boot):
+        idx = np.random.choice(n, size=n, replace=True)
+        diffs.append(np.mean(a[idx]) - np.mean(b[idx]))
+    diffs = np.array(diffs)
+    # two-sided p-value
+    p = float(np.mean(np.abs(diffs) >= np.abs(obs_diff)))
+    return p
+
+
+# ── Main ───────────────────────────────────────────────────────────────────────
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--n-boot', type=int, default=10000)
+    parser.add_argument('--out', default=os.path.join(DATA_DIR, 'compare_results.json'))
+    args_cli = parser.parse_args()
+
+    random.seed(42); np.random.seed(42); torch.manual_seed(42)
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f'[CMP] Device: {device}')
+
+    print('[CMP] Loading TabuLM...')
+    tabulm_model, tabulm_args, kv, bpe, enc_fn = load_tabulm(device)
+    print('[CMP] Loading KinyaBERT...')
+    kb_model, kb_tok, kb_ser = load_kinyabert(device)
+
+    with open(TABQA_FILE) as f:
+        all_items = json.load(f)
+    random.shuffle(all_items)
+    dev_items = all_items[int(0.8 * len(all_items)):]
+    print(f'[CMP] {len(dev_items)} dev items')
+
+    records = []
+    for item in dev_items:
+        csv_path = os.path.join(CSV_DIR, item['table_file'])
+        atype    = item.get('answer_type', '?')
+        if atype == 'count' or not os.path.exists(csv_path):
+            continue
+
+        cells = serialize_csv(csv_path)
+        if not cells:
+            continue
+
+        gold_rc = find_gold_cell(cells, item['answer'], item['question'])
+        if gold_rc is None:
+            continue
+
+        # TabuLM prediction
+        t_pred = tabulm_predict(tabulm_model, tabulm_args, kv, bpe, enc_fn,
+                                cells, item['question'], atype, device)
+        t_hit  = int(t_pred == gold_rc) if t_pred is not None else None
+
+        # KinyaBERT prediction
+        k_pred = kb_predict(kb_model, kb_tok, kb_ser,
+                            cells, item['question'], atype, device)
+        k_hit  = int(k_pred == gold_rc) if k_pred is not None else None
+
+        # only count items both models can evaluate
+        if t_hit is None or k_hit is None:
+            continue
+
+        # get cell text for readable output
+        cell_map = {(c.row_id, c.col_id): c.content.strip() for c in cells}
+        records.append({
+            'question':    item['question'],
+            'table':       item['table_file'],
+            'answer_type': atype,
+            'gold':        item['answer'],
+            'gold_rc':     list(gold_rc),
+            'tabulm_rc':   list(t_pred),
+            'tabulm_pred': cell_map.get(t_pred, '?'),
+            'tabulm_hit':  t_hit,
+            'kb_rc':       list(k_pred),
+            'kb_pred':     cell_map.get(k_pred, '?'),
+            'kb_hit':      k_hit,
+        })
+
+    print(f'\n[CMP] Paired items: {len(records)}')
+
+    # ── Per-type breakdown ──────────────────────────────────────────────────
+    by_type: Dict[str, Dict] = {}
+    for r in records:
+        t = r['answer_type']
+        by_type.setdefault(t, {'tabulm': [], 'kb': []})
+        by_type[t]['tabulm'].append(r['tabulm_hit'])
+        by_type[t]['kb'].append(r['kb_hit'])
+
+    print(f'\n{"Type":<14} {"N":>4}  {"TabuLM":>8}  {"KinyaBERT":>10}')
+    print('-' * 44)
+    all_t, all_k = [], []
+    for t, d in sorted(by_type.items()):
+        n = len(d['tabulm'])
+        t_em = sum(d['tabulm']) / n
+        k_em = sum(d['kb']) / n
+        all_t.extend(d['tabulm']); all_k.extend(d['kb'])
+        print(f'{t:<14} {n:>4}  {t_em:>8.3f}  {k_em:>10.3f}')
+    n_all = len(all_t)
+    print(f'{"Overall":<14} {n_all:>4}  {sum(all_t)/n_all:>8.3f}  {sum(all_k)/n_all:>10.3f}')
+
+    # ── Bootstrap CI ───────────────────────────────────────────────────────
+    print(f'\n[CMP] Bootstrap CI (n_boot={args_cli.n_boot}):')
+    t_em, t_lo, t_hi = bootstrap_ci(all_t, args_cli.n_boot)
+    k_em, k_lo, k_hi = bootstrap_ci(all_k, args_cli.n_boot)
+    p_val = paired_bootstrap_pvalue(all_t, all_k, args_cli.n_boot)
+    print(f'  TabuLM:    {t_em:.3f}  95% CI [{t_lo:.3f}, {t_hi:.3f}]')
+    print(f'  KinyaBERT: {k_em:.3f}  95% CI [{k_lo:.3f}, {k_hi:.3f}]')
+    print(f'  Two-sided paired bootstrap p-value: {p_val:.4f}')
+
+    # ── Error analysis: TabuLM correct, KinyaBERT wrong ────────────────────
+    wins = [r for r in records if r['tabulm_hit'] == 1 and r['kb_hit'] == 0]
+    losses = [r for r in records if r['tabulm_hit'] == 0 and r['kb_hit'] == 1]
+    print(f'\n[CMP] TabuLM wins (T=1, KB=0): {len(wins)}')
+    print(f'[CMP] TabuLM losses (T=0, KB=1): {len(losses)}')
+
+    print('\n=== TabuLM WINS (comparison type preferred) ===')
+    cmp_wins = [r for r in wins if r['answer_type'] == 'comparison']
+    for r in cmp_wins[:5]:
+        print(f'\n  Table : {r["table"]}')
+        print(f'  Q     : {r["question"]}')
+        print(f'  Type  : {r["answer_type"]}')
+        print(f'  Gold  : {r["gold"]}  (row={r["gold_rc"][0]}, col={r["gold_rc"][1]})')
+        print(f'  TabuLM: {r["tabulm_pred"]}  ✓')
+        print(f'  KinyaB: {r["kb_pred"]}  ✗')
+
+    print('\n=== TabuLM LOSSES ===')
+    for r in losses[:3]:
+        print(f'\n  Table : {r["table"]}')
+        print(f'  Q     : {r["question"]}')
+        print(f'  Type  : {r["answer_type"]}')
+        print(f'  Gold  : {r["gold"]}')
+        print(f'  TabuLM: {r["tabulm_pred"]}  ✗')
+        print(f'  KinyaB: {r["kb_pred"]}  ✓')
+
+    # ── Save ───────────────────────────────────────────────────────────────
+    out_data = {
+        'n_paired':       n_all,
+        'tabulm_em':      round(t_em, 4),
+        'kb_em':          round(k_em, 4),
+        'tabulm_ci':      [round(t_lo, 4), round(t_hi, 4)],
+        'kb_ci':          [round(k_lo, 4), round(k_hi, 4)],
+        'p_value':        round(p_val, 4),
+        'by_type':        {t: {
+                               'n':           len(d['tabulm']),
+                               'tabulm_em':   round(sum(d['tabulm'])/len(d['tabulm']), 4),
+                               'kb_em':       round(sum(d['kb'])/len(d['kb']), 4),
+                           } for t, d in by_type.items()},
+        'tabulm_wins':    wins,
+        'tabulm_losses':  losses,
+        'records':        records,
+    }
+    with open(args_cli.out, 'w') as f:
+        json.dump(out_data, f, indent=2, ensure_ascii=False)
+    print(f'\n[CMP] Saved to {args_cli.out}')
+
+
+if __name__ == '__main__':
+    main()

code/eval_llm_agg_fewshot.py

@@ -0,0 +1,141 @@
+#!/usr/bin/env python3
+"""3-shot GPT-4o experiment on aggregation items only.
+Tests whether few-shot prompting closes the aggregation gap vs TabuLM.
+
+Usage:
+  python eval_llm_agg_fewshot.py --api-key sk-...
+"""
+
+import argparse
+import csv
+import json
+import os
+import random
+import sys
+import time
+from typing import List, Optional
+
+import requests
+
+DATA_DIR   = '/shared/scratch/0/tmp/v_ireddi_rakshitha_results/tabulm/data'
+CODE_DIR   = '/shared/scratch/0/tmp/v_ireddi_rakshitha_results/tabulm/code'
+TABQA_FILE = os.path.join(DATA_DIR, 'tabqa_kin.json')
+CSV_DIR    = os.path.join(DATA_DIR, 'tables')
+OPENAI_URL = "https://api.openai.com/v1/chat/completions"
+
+sys.path.insert(0, CODE_DIR)
+from tabular_serializer import serialize_csv, TableCell
+from eval_llm_baseline import csv_to_markdown, find_gold_cell, em_match
+
+
+FEW_SHOT_EXAMPLES = [
+    {
+        "table": "| Inzego | Umubare_w_Abaturage | Umubare_w_Inzu |\n|---|---|---|\n| Gasabo | 542000 | 89000 |\n| Kicukiro | 312000 | 54000 |\n| Nyarugenge | 274000 | 48000 |",
+        "question": "Ni iki gifite Umubare_w_Abaturage nyinshi?",
+        "answer": "Gasabo",
+        "explanation": "Gasabo has the highest Umubare_w_Abaturage (542000).",
+    },
+    {
+        "table": "| Igihingwa | Umusaruro_Toni | Akarere |\n|---|---|---|\n| Ibigori | 45200 | Musanze |\n| Ibishyimbo | 12300 | Musanze |\n| Ibirayi | 89000 | Musanze |",
+        "question": "Ni iki gifite Umusaruro_Toni nke cyane?",
+        "answer": "Ibishyimbo",
+        "explanation": "Ibishyimbo has the lowest Umusaruro_Toni (12300).",
+    },
+    {
+        "table": "| Umwaka | Umusaruro_Miliyari_Frw | Impinduka_% |\n|---|---|---|\n| 2020 | 1100.2 | +3.1 |\n| 2021 | 1189.5 | +8.1 |\n| 2022 | 1209.3 | +1.7 |",
+        "question": "Ni iki gifite Impinduka_% ntoya muri bose?",
+        "answer": "2022",
+        "explanation": "2022 has the lowest Impinduka_% (+1.7).",
+    },
+]
+
+
+def make_fewshot_prompt(table_md: str, question: str) -> str:
+    examples = ""
+    for ex in FEW_SHOT_EXAMPLES:
+        examples += (
+            f"Table:\n{ex['table']}\n\n"
+            f"Question: {ex['question']}\n"
+            f"Answer: {ex['answer']}\n\n"
+        )
+    return (
+        f"You are answering questions about data tables written in Kinyarwanda.\n"
+        f"The answer is always the NAME of an entity from the first column — NOT a number.\n"
+        f"Identify which entity has the highest or lowest value in the relevant column.\n\n"
+        f"Examples:\n\n{examples}"
+        f"Now answer the following:\n\n"
+        f"Table:\n{table_md}\n\n"
+        f"Question: {question}\n"
+        f"Answer:"
+    )
+
+
+def call_openai(prompt: str, api_key: str, retries: int = 3) -> Optional[str]:
+    for attempt in range(retries):
+        try:
+            headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
+            payload = {"model": "gpt-4o", "max_tokens": 64, "temperature": 0,
+                       "messages": [{"role": "user", "content": prompt}]}
+            r = requests.post(OPENAI_URL, headers=headers, json=payload, timeout=30)
+            r.raise_for_status()
+            return r.json()["choices"][0]["message"]["content"].strip()
+        except Exception as e:
+            if attempt < retries - 1:
+                time.sleep(2 ** attempt)
+            else:
+                print(f"  [API error] {e}")
+    return None
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--api-key', required=True)
+    parser.add_argument('--out', default=os.path.join(DATA_DIR, 'llm_agg_fewshot_results.json'))
+    args = parser.parse_args()
+
+    random.seed(42)
+    with open(TABQA_FILE) as f:
+        all_items = json.load(f)
+    random.shuffle(all_items)
+    dev_items = [x for x in all_items[int(0.8 * len(all_items)):]
+                 if x.get('answer_type') == 'aggregation']
+    print(f'[FewShot-Agg] {len(dev_items)} aggregation dev items')
+
+    correct, total, skipped = 0, 0, 0
+    records = []
+
+    for i, item in enumerate(dev_items):
+        csv_path = os.path.join(CSV_DIR, item['table_file'])
+        if not os.path.exists(csv_path):
+            skipped += 1; continue
+        cells = serialize_csv(csv_path)
+        if not cells:
+            skipped += 1; continue
+        gold_rc = find_gold_cell(cells, item['answer'], item['question'])
+        if gold_rc is None:
+            skipped += 1; continue
+        table_md = csv_to_markdown(csv_path)
+        if not table_md:
+            skipped += 1; continue
+
+        prompt = make_fewshot_prompt(table_md, item['question'])
+        pred   = call_openai(prompt, args.api_key)
+        if pred is None:
+            skipped += 1; continue
+
+        hit = int(em_match(pred, item['answer']))
+        correct += hit; total += 1
+        records.append({'question': item['question'], 'gold': item['answer'],
+                        'pred': pred, 'hit': hit})
+        print(f'  [{i+1}/{len(dev_items)}] gold={item["answer"]!r:20s} pred={pred!r:20s} hit={hit}')
+
+    em = correct / total if total else 0.0
+    print(f'\n[FewShot-Agg] EM = {em:.4f} ({correct}/{total}, {skipped} skipped)')
+    with open(args.out, 'w') as f:
+        json.dump({'em': round(em, 4), 'correct': correct, 'total': total,
+                   'skipped': skipped, 'records': records}, f, indent=2, ensure_ascii=False)
+    print(f'[FewShot-Agg] Saved to {args.out}')
+
+
+if __name__ == '__main__':
+    main()

code/eval_llm_baseline.py

@@ -0,0 +1,261 @@
+#!/usr/bin/env python3
+"""Zero-shot LLM baseline on TabQA-kin dev set. Supports OpenAI and Anthropic.
+
+Usage:
+  # OpenAI (GPT-4o)
+  python eval_llm_baseline.py --provider openai --api-key sk-...
+  python eval_llm_baseline.py --provider openai --api-key sk-... --model gpt-4o-mini
+
+  # Anthropic (Claude Haiku)
+  python eval_llm_baseline.py --provider anthropic --api-key sk-ant-...
+"""
+
+import argparse
+import csv
+import json
+import os
+import random
+import sys
+import time
+from typing import Dict, List, Optional, Tuple
+
+import requests
+
+DATA_DIR   = '/shared/scratch/0/tmp/v_ireddi_rakshitha_results/tabulm/data'
+CODE_DIR   = '/shared/scratch/0/tmp/v_ireddi_rakshitha_results/tabulm/code'
+TABQA_FILE = os.path.join(DATA_DIR, 'tabqa_kin.json')
+CSV_DIR    = os.path.join(DATA_DIR, 'tables')
+
+sys.path.insert(0, CODE_DIR)
+from tabular_serializer import serialize_csv, TableCell
+
+ANTHROPIC_URL = "https://api.anthropic.com/v1/messages"
+OPENAI_URL    = "https://api.openai.com/v1/chat/completions"
+
+
+# ── Table formatting ───────────────────────────────────────────────────────────
+
+def csv_to_markdown(csv_path: str) -> Optional[str]:
+    try:
+        with open(csv_path, encoding='utf-8') as f:
+            rows = list(csv.reader(f))
+        if not rows:
+            return None
+        lines = [' | '.join(r) for r in rows]
+        # insert separator after header
+        lines.insert(1, ' | '.join(['---'] * len(rows[0])))
+        return '\n'.join(lines)
+    except Exception:
+        return None
+
+
+# ── Gold cell lookup (same logic as finetune_tabqa.py v3) ─────────────────────
+
+def find_gold_cell(cells: List[TableCell], answer_text: str,
+                   question_text: str = '') -> Optional[Tuple[int, int]]:
+    answer_norm = answer_text.strip().lower()
+    matches = [(c.row_id, c.col_id) for c in cells
+               if c.row_id > 1 and c.col_id > 0
+               and c.content.strip() == answer_text.strip()]
+    if not matches:
+        matches = [(c.row_id, c.col_id) for c in cells
+                   if c.row_id > 1 and c.col_id > 0
+                   and c.content.strip().lower() == answer_norm]
+    if not matches:
+        return None
+    if len(matches) == 1:
+        return matches[0]
+    if question_text:
+        q_words = set(question_text.lower().split())
+        row_labels  = {c.row_id: c.content.strip().lower()
+                       for c in cells if c.col_id == 1 and c.row_id > 1}
+        col_headers = {c.col_id: c.content.strip().lower()
+                       for c in cells if c.row_id == 1 and c.col_id > 0}
+        best, best_score = None, (-1, -1)
+        for (row_id, col_id) in matches:
+            r_sc = len(q_words & set(row_labels.get(row_id,  '').split()))
+            c_sc = len(q_words & set(col_headers.get(col_id, '').split()))
+            if (r_sc, c_sc) > best_score:
+                best_score, best = (r_sc, c_sc), (row_id, col_id)
+        if best_score[0] > 0 or best_score[1] > 0:
+            return best
+    return min(matches, key=lambda rc: (rc[0], rc[1]))
+
+
+# ── Prompt construction ────────────────────────────────────────────────────────
+
+TYPE_HINTS = {
+    'lookup':      'The answer is a specific value (number or text) from the table. Find the row matching the entity named in the question, then return the value in the relevant column.',
+    'comparison':  'The answer is the NAME of one of the two entities mentioned in the question — NOT a number. Find both entities in the first column, compare their values in the specified data column, and return ONLY the entity name that has the higher or lower value as asked.',
+    'aggregation': 'The answer is the NAME of an entity from the first column of the table — NOT a number. Identify which entity has the highest or lowest value in the relevant column, then return that entity\'s name exactly as it appears in the table.',
+}
+
+def make_prompt(table_md: str, question: str, answer_type: str) -> str:
+    hint = TYPE_HINTS.get(answer_type, '')
+    return (
+        f"You are answering a question about a data table written in Kinyarwanda.\n\n"
+        f"Table:\n{table_md}\n\n"
+        f"Question: {question}\n\n"
+        f"{hint}\n\n"
+        f"Reply with ONLY the exact value from the table that answers the question. "
+        f"No explanation, no added words."
+    )
+
+
+# ── API call ───────────────────────────────────────────────────────────────────
+
+def call_llm(prompt: str, api_key: str, model: str, provider: str,
+             retries: int = 3) -> Optional[str]:
+    for attempt in range(retries):
+        try:
+            if provider == 'openai':
+                headers = {
+                    "Authorization": f"Bearer {api_key}",
+                    "Content-Type":  "application/json",
+                }
+                payload = {
+                    "model":      model,
+                    "max_tokens": 64,
+                    "messages":   [{"role": "user", "content": prompt}],
+                    "temperature": 0,
+                }
+                r = requests.post(OPENAI_URL, headers=headers,
+                                  json=payload, timeout=30)
+                r.raise_for_status()
+                return r.json()["choices"][0]["message"]["content"].strip()
+            else:  # anthropic
+                headers = {
+                    "x-api-key":         api_key,
+                    "anthropic-version": "2023-06-01",
+                    "content-type":      "application/json",
+                }
+                payload = {
+                    "model":      model,
+                    "max_tokens": 64,
+                    "messages":   [{"role": "user", "content": prompt}],
+                }
+                r = requests.post(ANTHROPIC_URL, headers=headers,
+                                  json=payload, timeout=30)
+                if r.status_code == 529:
+                    time.sleep(10 * (attempt + 1))
+                    continue
+                r.raise_for_status()
+                return r.json()["content"][0]["text"].strip()
+        except Exception as e:
+            if attempt < retries - 1:
+                time.sleep(2 ** attempt)
+            else:
+                print(f"  [API error] {e}")
+    return None
+
+
+# ── EM with normalization ──────────────────────────────────────────────────────
+
+def em_match(pred: str, gold: str) -> bool:
+    def norm(s: str) -> str:
+        s = s.strip().lower().replace(',', '').replace(' ', '')
+        try:
+            return f'{float(s):g}'
+        except ValueError:
+            return s
+    return norm(pred) == norm(gold)
+
+
+# ── Main ───────────────────────────────────────────────────────────────────────
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--api-key',  required=True, help='OpenAI or Anthropic API key')
+    parser.add_argument('--provider', default='openai', choices=['openai', 'anthropic'])
+    parser.add_argument('--model',    default=None,
+                        help='Model name (default: gpt-4o for OpenAI, claude-haiku-4-5-20251001 for Anthropic)')
+    parser.add_argument('--out',      default=os.path.join(DATA_DIR, 'llm_baseline_results.json'))
+    parser.add_argument('--max-items', type=int, default=None,
+                        help='Limit number of dev items (for testing)')
+    args = parser.parse_args()
+    if args.model is None:
+        args.model = 'gpt-4o' if args.provider == 'openai' else 'claude-haiku-4-5-20251001'
+
+    random.seed(42)
+    with open(TABQA_FILE) as f:
+        all_items = json.load(f)
+    random.shuffle(all_items)
+    dev_items = all_items[int(0.8 * len(all_items)):]
+    if args.max_items:
+        dev_items = dev_items[:args.max_items]
+    print(f'[LLM] Provider: {args.provider}  Model: {args.model}')
+    print(f'[LLM] {len(dev_items)} dev items')
+
+    correct, total, skipped = 0, 0, 0
+    by_type: Dict[str, List[int]] = {}
+    records = []
+
+    for i, item in enumerate(dev_items):
+        csv_path = os.path.join(CSV_DIR, item['table_file'])
+        atype    = item.get('answer_type', '?')
+
+        # skip count questions (answer not a cell value)
+        if atype == 'count':
+            skipped += 1; continue
+
+        if not os.path.exists(csv_path):
+            skipped += 1; continue
+
+        cells = serialize_csv(csv_path)
+        if not cells:
+            skipped += 1; continue
+
+        # verify gold cell exists (same filter as our fine-tuning eval)
+        gold_rc = find_gold_cell(cells, item['answer'], item['question'])
+        if gold_rc is None:
+            skipped += 1; continue
+
+        table_md = csv_to_markdown(csv_path)
+        if not table_md:
+            skipped += 1; continue
+
+        prompt = make_prompt(table_md, item['question'], atype)
+        pred   = call_llm(prompt, args.api_key, args.model, args.provider)
+        if pred is None:
+            skipped += 1; continue
+
+        hit = int(em_match(pred, item['answer']))
+        correct += hit
+        total   += 1
+        by_type.setdefault(atype, []).append(hit)
+        records.append({
+            'idx':          i,
+            'question':     item['question'],
+            'table':        item['table_file'],
+            'answer_type':  atype,
+            'gold':         item['answer'],
+            'pred':         pred,
+            'hit':          hit,
+        })
+
+        if (i + 1) % 10 == 0:
+            running_em = correct / total if total else 0.0
+            print(f'  [{i+1:3d}/{len(dev_items)}] running EM={running_em:.3f}  '
+                  f'correct={correct}/{total}  skipped={skipped}')
+
+    em = correct / total if total else 0.0
+    print(f'\n[LLM] Final EM = {em:.4f} ({correct}/{total}, {skipped} skipped)')
+    for t, hits in sorted(by_type.items()):
+        print(f'  {t}: {sum(hits)}/{len(hits)} = {sum(hits)/len(hits):.3f}')
+
+    result = {
+        'model':    args.model,
+        'em':       round(em, 4),
+        'correct':  correct,
+        'total':    total,
+        'skipped':  skipped,
+        'by_type':  {t: round(sum(h)/len(h), 4) for t, h in by_type.items()},
+        'records':  records,
+    }
+    with open(args.out, 'w') as f:
+        json.dump(result, f, indent=2, ensure_ascii=False)
+    print(f'[LLM] Results saved to {args.out}')
+
+
+if __name__ == '__main__':
+    main()

code/finetune_tableqa.py

@@ -0,0 +1,244 @@
+"""
+finetune_tableqa.py — Fine-tune TabuLM on TabQA-kin benchmark.
+
+Fine-tuning approach:
+  - Serialize table + question → token string via tabular_serializer
+  - Append question tokens after [SEP]
+  - Predict answer as span (start/end token indices) for lookup/comparison/aggregation
+  - Predict integer for count questions
+  - Metrics: exact match (EM) and token-level F1
+
+Usage:
+    python code/finetune_tableqa.py \
+        --checkpoint checkpoints/tabulm_step120000.pt \
+        --tabqa     data/tabqa_kin.json \
+        --tables    data/tables/ \
+        --output    checkpoints/tabulm_tableqa/
+"""
+
+import argparse
+import json
+import os
+import random
+from pathlib import Path
+from typing import Dict, List, Tuple
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, Dataset
+
+random.seed(42)
+torch.manual_seed(42)
+
+# ── Span prediction head ──────────────────────────────────────────────────────
+
+class SpanHead(nn.Module):
+    def __init__(self, hidden_size: int):
+        super().__init__()
+        self.start_linear = nn.Linear(hidden_size, 1)
+        self.end_linear   = nn.Linear(hidden_size, 1)
+
+    def forward(self, hidden_states):
+        start_logits = self.start_linear(hidden_states).squeeze(-1)
+        end_logits   = self.end_linear(hidden_states).squeeze(-1)
+        return start_logits, end_logits
+
+
+class TabuLM_QA(nn.Module):
+    def __init__(self, encoder, hidden_size: int = 768):
+        super().__init__()
+        self.encoder   = encoder
+        self.span_head = SpanHead(hidden_size)
+        self.count_head = nn.Linear(hidden_size, 100)
+
+    def forward(self, batch, answer_type: str = 'lookup'):
+        hidden = self.encoder(batch)['last_hidden_state']
+        if answer_type == 'count':
+            cls_hidden = hidden[:, 0, :]
+            return self.count_head(cls_hidden)
+        else:
+            return self.span_head(hidden)
+
+
+# ── Dataset ───────────────────────────────────────────────────────────────────
+
+class TabQADataset(Dataset):
+    def __init__(self, items: List[Dict], tables_dir: Path, kb_vocab, bpe):
+        from tabular_serializer import serialize_csv, table_cells_to_text
+        from morpho_stub import parse_text_stub
+
+        self.samples = []
+        for item in items:
+            table_path = tables_dir / item['table_file']
+            if not table_path.exists():
+                continue
+            cells = serialize_csv(str(table_path))
+            table_text, word_meta = table_cells_to_text(cells)
+
+            question = item['question']
+            full_text = table_text + ' [SEP] ' + question
+
+            tokens = parse_text_stub(full_text, kb_vocab, bpe)
+            self.samples.append({
+                'tokens':      tokens,
+                'word_meta':   word_meta,
+                'answer':      item['answer'],
+                'answer_type': item['answer_type'],
+                'answer_row':  item.get('answer_row', -1),
+                'answer_col':  item.get('answer_col', -1),
+                'id':          item['id'],
+            })
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        return self.samples[idx]
+
+
+# ── Metrics ───────────────────────────────────────────────────────────────────
+
+def exact_match(pred: str, gold: str) -> float:
+    return float(pred.strip().lower() == gold.strip().lower())
+
+
+def token_f1(pred: str, gold: str) -> float:
+    pred_toks = pred.strip().lower().split()
+    gold_toks = gold.strip().lower().split()
+    common = set(pred_toks) & set(gold_toks)
+    if not common:
+        return 0.0
+    p = len(common) / len(pred_toks) if pred_toks else 0
+    r = len(common) / len(gold_toks) if gold_toks else 0
+    if p + r == 0:
+        return 0.0
+    return 2 * p * r / (p + r)
+
+
+def evaluate(model, dataloader, device) -> Dict[str, float]:
+    model.eval()
+    em_sum = f1_sum = total = 0
+    with torch.no_grad():
+        for batch in dataloader:
+            answer_type = batch['answer_type'][0]
+            preds = model(batch, answer_type=answer_type)
+
+            if answer_type == 'count':
+                pred_counts = preds.argmax(dim=-1).tolist()
+                for pred, gold in zip(pred_counts, batch['answer']):
+                    em_sum  += exact_match(str(pred), gold)
+                    f1_sum  += token_f1(str(pred), gold)
+                    total   += 1
+            else:
+                start_logits, end_logits = preds
+                for i in range(len(batch['answer'])):
+                    s = start_logits[i].argmax().item()
+                    e = end_logits[i].argmax().item()
+                    tokens = batch['tokens'][i]
+                    pred_text = ' '.join(t.stem for t in tokens[s:e+1]) if s <= e else ''
+                    gold_text = batch['answer'][i]
+                    em_sum  += exact_match(pred_text, gold_text)
+                    f1_sum  += token_f1(pred_text, gold_text)
+                    total   += 1
+
+    return {
+        'exact_match': em_sum / total if total else 0,
+        'f1':          f1_sum / total if total else 0,
+        'total':       total,
+    }
+
+
+# ── Training loop ─────────────────────────────────────────────────────────────
+
+def train(args):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f'Device: {device}')
+
+    with open(args.tabqa, encoding='utf-8') as f:
+        all_items = json.load(f)
+
+    random.shuffle(all_items)
+    split = int(0.8 * len(all_items))
+    train_items = all_items[:split]
+    dev_items   = all_items[split:]
+    print(f'Train: {len(train_items)}  Dev: {len(dev_items)}')
+
+    import sys; sys.path.insert(0, 'code')
+    from morpho_common import KBVocab
+    from morpho_stub   import BPEFallback
+
+    kb_vocab = KBVocab.load(args.kb_vocab)
+    bpe      = BPEFallback(args.bpe_codes)
+
+    train_ds = TabQADataset(train_items, Path(args.tables), kb_vocab, bpe)
+    dev_ds   = TabQADataset(dev_items,   Path(args.tables), kb_vocab, bpe)
+
+    checkpoint = torch.load(args.checkpoint, map_location='cpu')
+    from tabulm_model import tabulm_base
+    model_core = tabulm_base(kb_vocab)
+    model_core.load_state_dict(checkpoint['model'], strict=False)
+
+    model = TabuLM_QA(model_core.encoder).to(device)
+
+    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.01)
+    ce_loss   = nn.CrossEntropyLoss()
+
+    Path(args.output).mkdir(parents=True, exist_ok=True)
+    best_em = 0.0
+
+    for epoch in range(args.epochs):
+        model.train()
+        total_loss = 0
+        for step, batch in enumerate(DataLoader(train_ds, batch_size=args.batch_size, shuffle=True)):
+            optimizer.zero_grad()
+            answer_type = batch['answer_type'][0]
+            preds = model(batch, answer_type=answer_type)
+
+            if answer_type == 'count':
+                labels = torch.tensor([int(a) for a in batch['answer']], device=device)
+                loss = ce_loss(preds, labels)
+            else:
+                start_logits, end_logits = preds
+                start_labels = torch.tensor(batch['answer_row'], device=device)
+                end_labels   = torch.tensor(batch['answer_row'], device=device)
+                loss = ce_loss(start_logits, start_labels) + ce_loss(end_logits, end_labels)
+
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            optimizer.step()
+            total_loss += loss.item()
+
+            if step % 50 == 0:
+                print(f'  epoch {epoch+1} step {step}  loss={total_loss/(step+1):.4f}')
+
+        metrics = evaluate(model, DataLoader(dev_ds, batch_size=32), device)
+        print(f'Epoch {epoch+1}  EM={metrics["exact_match"]:.4f}  F1={metrics["f1"]:.4f}')
+
+        if metrics['exact_match'] > best_em:
+            best_em = metrics['exact_match']
+            ckpt_path = Path(args.output) / 'best_tableqa.pt'
+            torch.save({'model': model.state_dict(), 'metrics': metrics}, ckpt_path)
+            print(f'  → saved best checkpoint  EM={best_em:.4f}')
+
+    print(f'\nBest EM: {best_em:.4f}')
+
+
+# ── Entry point ───────────────────────────────────────────────────────────────
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--checkpoint', required=True)
+    parser.add_argument('--tabqa',      default='data/tabqa_kin.json')
+    parser.add_argument('--tables',     default='data/tables')
+    parser.add_argument('--output',     default='checkpoints/tabulm_tableqa')
+    parser.add_argument('--kb_vocab',   default='conf/kb_vocab_state_dict_2021-02-07.pt')
+    parser.add_argument('--bpe_codes',  default='conf/bpe_codes.txt')
+    parser.add_argument('--epochs',     type=int,   default=5)
+    parser.add_argument('--batch_size', type=int,   default=16)
+    parser.add_argument('--lr',         type=float, default=2e-5)
+    args = parser.parse_args()
+    train(args)
+
+
+if __name__ == '__main__':
+    main()

code/finetune_tabqa.py

@@ -0,0 +1,641 @@
+#!/usr/bin/env python3
+"""Fine-tune TabuLM for cell-selection on TabQA-kin (526 Kinyarwanda table QA pairs).
+
+Strategy: Given a (table, question) pair, encode both through the TabuLM encoder
+and predict which table cell (row_id, col_id) contains the answer via a linear
+scoring head over per-cell pooled hidden states.
+"""
+
+import argparse
+import json
+import os
+import random
+import sys
+from datetime import datetime
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn.utils import clip_grad_norm_
+
+# ── Server paths ───────────────────────────────────────────────────────────────
+CODE_DIR = '/shared/scratch/0/tmp/v_ireddi_rakshitha_results/tabulm/code'
+DATA_DIR = '/shared/scratch/0/tmp/v_ireddi_rakshitha_results/tabulm/data'
+CHECKPOINT = os.path.join(DATA_DIR, 'tabulm_model_2026-05-13_pos@1_stem@1_afsets@False@10000.pt')
+TABQA_FILE = os.path.join(DATA_DIR, 'tabqa_kin.json')
+CSV_DIR    = os.path.join(DATA_DIR, 'tables')
+BEST_MODEL = os.path.join(DATA_DIR, 'finetune_tabqa_best.pt')
+RESULTS    = os.path.join(DATA_DIR, 'finetune_tabqa_results.json')
+
+sys.path.insert(0, CODE_DIR)
+
+import youtokentome as yttm
+from morpho_data_loaders import KBVocab
+from tabular_serializer import serialize_csv, table_cells_to_text, TableCell
+from morpho_stub import parse_text_stub
+from tabulm_model import TabuLM, tabulm_base
+
+
+# ── Architecture auto-detection ────────────────────────────────────────────────
+
+def detect_arch_from_state(state: dict) -> argparse.Namespace:
+    """Infer all architecture hyperparameters by inspecting checkpoint tensor shapes."""
+    ns = argparse.Namespace()
+
+    # Sequence transformer model dimension
+    ns.seq_tr_d_model = int(state['encoder.row_embedding.weight'].shape[1])
+    d = ns.seq_tr_d_model
+
+    # Stem / morpho dims
+    ns.stem_dim  = int(state['encoder.s_stem_embedding.weight'].shape[1])
+    if 'encoder.m1_pos_embedding.weight' in state:
+        ns.morpho_dim = int(state['encoder.m1_pos_embedding.weight'].shape[1])
+    elif 'encoder.m_stem_embedding.weight' in state:
+        ns.morpho_dim = int(state['encoder.m_stem_embedding.weight'].shape[1])
+    else:
+        ns.morpho_dim = 128
+
+    # Sequence transformer layers
+    layer_ids = {
+        int(k.split('.layers.')[1].split('.')[0])
+        for k in state
+        if 'encoder.seq_transformer_encoder.layers.' in k and '.layers.' in k
+    }
+    ns.seq_tr_nlayers = max(layer_ids) + 1 if layer_ids else 12
+
+    # Heads from tabular bias parameter
+    if 'encoder.row_attn_bias' in state:
+        ns.seq_tr_nhead = int(state['encoder.row_attn_bias'].shape[0])
+    else:
+        ns.seq_tr_nhead = 8
+
+    # Feed-forward dim
+    ff_key = 'encoder.seq_transformer_encoder.layers.0.linear1.weight'
+    ns.seq_tr_dim_feedforward = int(state[ff_key].shape[0]) if ff_key in state else d * 4
+
+    # Morpho transformer layers
+    m_ids = {
+        int(k.split('.layers.')[1].split('.')[0])
+        for k in state
+        if 'encoder.morpho_transformer_encoder.layers.' in k and '.layers.' in k
+    }
+    ns.morpho_tr_nlayers = max(m_ids) + 1 if m_ids else 4
+
+    mff_key = 'encoder.morpho_transformer_encoder.layers.0.linear1.weight'
+    ns.morpho_tr_dim_feedforward = int(state[mff_key].shape[0]) if mff_key in state else ns.morpho_dim * 4
+
+    ns.morpho_tr_nhead = max(1, ns.morpho_dim // 32)  # head_dim=32 default
+    ns.morpho_tr_dropout = 0.1
+    ns.seq_tr_dropout    = 0.1
+    ns.layernorm_epsilon = 1e-6
+    ns.max_seq_len = 512
+
+    # Tabular embedding flags
+    ns.num_pos_m_embeddings  = 1
+    ns.num_stem_m_embeddings = 1
+    ns.use_afsets            = False
+    ns.afset_dict_size       = 10000
+    ns.use_morpho_encoder    = True
+
+    morpho_part     = d - ns.stem_dim
+    tot_morpho_vecs = morpho_part // ns.morpho_dim if ns.morpho_dim else 2
+    tot_morpho_idx  = ns.num_pos_m_embeddings + ns.num_stem_m_embeddings
+    ns.use_affix_bow_m_embedding = (tot_morpho_vecs > tot_morpho_idx)
+
+    ns.use_tupe_rel_pos_bias       = 'encoder.tupe_rel_pos_bias.weight' in state \
+                                     or any('tupe' in k for k in state)
+    ns.use_pos_aware_rel_pos_bias  = any('pos_aware' in k for k in state)
+    ns.use_pos_aware_rel           = ns.use_pos_aware_rel_pos_bias
+    ns.predict_affixes             = False
+
+    # Needed by tabulm_base factory
+    ns.gpus            = 0
+    ns.world_size      = 1
+    ns.exploratory_model_load = None
+
+    return ns
+
+
+# ── Gold cell lookup via text matching ────────────────────────────────────────
+
+def find_gold_cell(cells: List[TableCell], answer_text: str,
+                   question_text: str = '') -> Optional[Tuple[int, int]]:
+    """
+    Find the (row_id, col_id) of the cell whose content matches answer_text.
+    When multiple cells share the same value, uses question-word overlap against
+    both first-column row labels AND header-row column names to disambiguate.
+    """
+    answer_norm = answer_text.strip().lower()
+
+    # Pass 1: exact match
+    matches = [(c.row_id, c.col_id) for c in cells
+               if c.row_id > 1 and c.col_id > 0 and c.content.strip() == answer_text.strip()]
+    if len(matches) == 1:
+        return matches[0]
+
+    # Pass 2: case-insensitive
+    if not matches:
+        matches = [(c.row_id, c.col_id) for c in cells
+                   if c.row_id > 1 and c.col_id > 0 and c.content.strip().lower() == answer_norm]
+
+    if not matches:
+        return None
+    if len(matches) == 1:
+        return matches[0]
+
+    # Multiple cells match — score by (row overlap, col overlap) jointly
+    if question_text:
+        q_words = set(question_text.lower().split())
+        row_labels  = {c.row_id: c.content.strip().lower()
+                       for c in cells if c.col_id == 1 and c.row_id > 1}
+        col_headers = {c.col_id: c.content.strip().lower()
+                       for c in cells if c.row_id == 1 and c.col_id > 0}
+        best, best_score = None, (-1, -1)
+        for (row_id, col_id) in matches:
+            r_sc = len(q_words & set(row_labels.get(row_id,  '').split()))
+            c_sc = len(q_words & set(col_headers.get(col_id, '').split()))
+            if (r_sc, c_sc) > best_score:
+                best_score, best = (r_sc, c_sc), (row_id, col_id)
+        if best_score[0] > 0 or best_score[1] > 0:
+            return best
+
+    return min(matches, key=lambda rc: (rc[0], rc[1]))
+
+
+# ── Sequence encoding (no masking) ────────────────────────────────────────────
+
+def _add_special(key, kv, pos_tags, stems, tokens_lengths, row_ids, col_ids, cell_types):
+    pos_tags.append(kv.pos_tag_vocab[key])
+    stems.append(kv.reduced_stem_vocab[key])
+    tokens_lengths.append(0)
+    row_ids.append(0)
+    col_ids.append(0)
+    cell_types.append(0)
+
+
+def encode_item(
+    cells,
+    question_text: str,
+    kv: KBVocab,
+    bpe,
+    max_seq_len: int = 512,
+) -> Optional[Tuple]:
+    """
+    Encode a table (List[TableCell]) + question string into flat lists for the
+    TabuLM encoder. Returns None if the table serialization fails.
+
+    Returns:
+      pos_tags, stems, affixes, tokens_lengths,
+      row_ids, col_ids, cell_types,
+      ordered_cells, cell_to_positions
+    where ordered_cells is a sorted list of unique (row_id, col_id) tuples, and
+    cell_to_positions maps each (row_id, col_id) → [token_idx, ...].
+    """
+    text, word_meta = table_cells_to_text(cells)
+    parsed_table = parse_text_stub(text, kv, bpe)
+    if len(parsed_table) != len(word_meta):
+        return None
+
+    parsed_question = parse_text_stub(question_text, kv, bpe)
+
+    pos_tags, stems, affixes, tokens_lengths = [], [], [], []
+    row_ids, col_ids, cell_types = [], [], []
+    cell_to_positions: Dict[Tuple[int, int], List[int]] = {}
+
+    # CLS
+    _add_special('<CLS>', kv, pos_tags, stems, tokens_lengths, row_ids, col_ids, cell_types)
+
+    # Table tokens — each ParsedToken may have multiple sub-stems (BPE)
+    for pt, (r, c, ct) in zip(parsed_table, word_meta):
+        for sidx in pt.stem_idx:
+            seq_idx = len(pos_tags)
+            pos_tags.append(pt.pos_tag_idx)
+            stems.append(kv.mapped_stem_vocab_idx[sidx])
+            affixes.extend(pt.affixes_idx)
+            tokens_lengths.append(len(pt.affixes_idx))
+            row_ids.append(r)
+            col_ids.append(c)
+            cell_types.append(ct)
+            if r > 0 and c > 0:
+                cell_to_positions.setdefault((r, c), []).append(seq_idx)
+
+    # SEP between table and question
+    _add_special('<SEP>', kv, pos_tags, stems, tokens_lengths, row_ids, col_ids, cell_types)
+
+    # Question tokens (row_id=0, col_id=0 → not scored as table cells)
+    for pt in parsed_question:
+        for sidx in pt.stem_idx:
+            pos_tags.append(pt.pos_tag_idx)
+            stems.append(kv.mapped_stem_vocab_idx[sidx])
+            affixes.extend(pt.affixes_idx)
+            tokens_lengths.append(len(pt.affixes_idx))
+            row_ids.append(0)
+            col_ids.append(0)
+            cell_types.append(0)
+
+    # Final SEP
+    _add_special('<SEP>', kv, pos_tags, stems, tokens_lengths, row_ids, col_ids, cell_types)
+
+    # Truncate to max_seq_len
+    if len(pos_tags) > max_seq_len:
+        pos_tags      = pos_tags[:max_seq_len]
+        stems         = stems[:max_seq_len]
+        tokens_lengths = tokens_lengths[:max_seq_len]
+        row_ids       = row_ids[:max_seq_len]
+        col_ids       = col_ids[:max_seq_len]
+        cell_types    = cell_types[:max_seq_len]
+        affixes       = affixes[:sum(tokens_lengths)]
+        cell_to_positions = {
+            k: [p for p in v if p < max_seq_len]
+            for k, v in cell_to_positions.items()
+        }
+        cell_to_positions = {k: v for k, v in cell_to_positions.items() if v}
+
+    ordered_cells = sorted(cell_to_positions.keys(),
+                           key=lambda rc: min(cell_to_positions[rc]))
+
+    return (pos_tags, stems, affixes, tokens_lengths,
+            row_ids, col_ids, cell_types,
+            ordered_cells, cell_to_positions)
+
+
+# ── Fine-tuning model ──────────────────────────────────────────────────────────
+
+class TabQAModel(nn.Module):
+    """TabuLM encoder + linear cell-selection head."""
+
+    def __init__(self, tabulm: TabuLM):
+        super().__init__()
+        self.encoder   = tabulm.encoder
+        d = self.encoder.seq_tr_d_model
+        self.cell_head = nn.Linear(d, 1)
+        nn.init.normal_(self.cell_head.weight, std=0.02)
+        nn.init.zeros_(self.cell_head.bias)
+
+    def get_hidden(self, args, pos_tags, stems, affixes,
+                   tokens_lengths, row_ids, col_ids, cell_types, device):
+        pos_t  = torch.tensor(pos_tags,     dtype=torch.long, device=device)
+        stem_t = torch.tensor(stems,        dtype=torch.long, device=device)
+        afx_t  = (torch.tensor(affixes, dtype=torch.long, device=device)
+                  if affixes else torch.zeros(0, dtype=torch.long, device=device))
+        row_t  = torch.tensor(row_ids,      dtype=torch.long, device=device)
+        col_t  = torch.tensor(col_ids,      dtype=torch.long, device=device)
+        ct_t   = torch.tensor(cell_types,   dtype=torch.long, device=device)
+
+        hidden = self.encoder.forward(
+            args,
+            rel_pos_arr=None,
+            tokens_lengths=tokens_lengths,
+            input_sequence_lengths=[len(pos_tags)],
+            pos_tags=pos_t, stems=stem_t, afsets=None, affixes=afx_t,
+            row_ids=row_t, col_ids=col_t, cell_types=ct_t,
+        )  # (S, 1, d)
+        return hidden[:, 0, :]  # (S, d)
+
+    def forward(self, args, pos_tags, stems, affixes,
+                tokens_lengths, row_ids, col_ids, cell_types,
+                ordered_cells, cell_to_positions, device):
+        """
+        Returns (scores, valid_cells) where scores is a (C,) tensor of logits
+        and valid_cells is the subset of ordered_cells that had token positions.
+        """
+        hidden = self.get_hidden(args, pos_tags, stems, affixes,
+                                 tokens_lengths, row_ids, col_ids, cell_types, device)
+        S = hidden.size(0)
+
+        cell_embeds, valid_cells = [], []
+        for rc in ordered_cells:
+            positions = [p for p in cell_to_positions[rc] if p < S]
+            if not positions:
+                continue
+            h = hidden[positions].mean(0)
+            cell_embeds.append(h)
+            valid_cells.append(rc)
+
+        if not cell_embeds:
+            return None, []
+
+        cell_embeds = torch.stack(cell_embeds)          # (C, d)
+        scores      = self.cell_head(cell_embeds).squeeze(-1)  # (C,)
+        return scores, valid_cells
+
+
+# ── Evaluation ────────────────────────────────────────────────────────────────
+
+def _predict_lookup(scores, valid_cells, cells, question_text):
+    """For lookup questions: restrict to (best_row, best_col) using question-word
+    overlap against row labels (first column) and column headers (header row)."""
+    q_words = set(question_text.lower().split())
+    row_labels  = {c.row_id: c.content.strip().lower()
+                   for c in cells if c.col_id == 1 and c.row_id > 1}
+    col_headers = {c.col_id: c.content.strip().lower()
+                   for c in cells if c.row_id == 1 and c.col_id > 0}
+
+    if not q_words:
+        return valid_cells[scores.argmax().item()]
+
+    row_score = {r: len(q_words & set(lbl.split())) for r, lbl in row_labels.items()}
+    col_score = {c: len(q_words & set(hdr.split())) for c, hdr in col_headers.items()}
+
+    best_row = max(row_score, key=row_score.get) if row_score else None
+    best_col = max(col_score, key=col_score.get) if col_score else None
+    has_row  = best_row is not None and row_score[best_row] > 0
+    has_col  = best_col is not None and col_score[best_col] > 0
+
+    if not has_row and not has_col:
+        return valid_cells[scores.argmax().item()]
+
+    # Narrow to best row first
+    if has_row:
+        row_indices = [i for i, rc in enumerate(valid_cells) if rc[0] == best_row]
+    else:
+        row_indices = list(range(len(valid_cells)))
+
+    if not row_indices:
+        return valid_cells[scores.argmax().item()]
+
+    # Further narrow to best column if signal available
+    if has_col:
+        both = [i for i in row_indices if valid_cells[i][1] == best_col]
+        if both:
+            return valid_cells[max(both, key=lambda i: scores[i].item())]
+
+    return valid_cells[max(row_indices, key=lambda i: scores[i].item())]
+
+
+def _predict_comparison(scores, valid_cells, cells, question_text):
+    """For comparison questions: restrict to col-1 (entity name) cells of the
+    top-2 question-relevant rows. The answer is the winning entity name."""
+    q_words = set(question_text.lower().split())
+    row_labels = {c.row_id: c.content.strip().lower()
+                  for c in cells if c.col_id == 1 and c.row_id > 1}
+
+    if not q_words or not row_labels:
+        return valid_cells[scores.argmax().item()]
+
+    row_score = {r: len(q_words & set(lbl.split())) for r, lbl in row_labels.items()}
+    top2 = sorted([r for r, s in row_score.items() if s > 0],
+                  key=lambda r: row_score[r], reverse=True)[:2]
+
+    if not top2:
+        return valid_cells[scores.argmax().item()]
+
+    # Prefer col-1 cells (entity names) in those rows; fall back to any cell
+    cand = [i for i, (r, c) in enumerate(valid_cells) if r in top2 and c == 1]
+    if not cand:
+        cand = [i for i, (r, c) in enumerate(valid_cells) if r in top2]
+    if not cand:
+        return valid_cells[scores.argmax().item()]
+
+    return valid_cells[max(cand, key=lambda i: scores[i].item())]
+
+
+def evaluate(model, args, items, kv, bpe, csv_dir, device, split='dev'):
+    model.eval()
+    correct  = 0
+    total    = 0
+    skipped  = 0
+    by_type: Dict[str, List[int]] = {}
+
+    with torch.no_grad():
+        for item in items:
+            csv_path = os.path.join(csv_dir, item['table_file'])
+            if not os.path.exists(csv_path):
+                skipped += 1
+                continue
+            cells = serialize_csv(csv_path)
+            if not cells:
+                skipped += 1
+                continue
+
+            question = item['question']
+            atype    = item.get('answer_type', 'unknown')
+
+            # Find gold cell — use question context to resolve duplicate values
+            gold_rc = find_gold_cell(cells, item['answer'], question_text=question)
+            if gold_rc is None:
+                skipped += 1
+                continue
+
+            enc = encode_item(cells, question, kv, bpe)
+            if enc is None:
+                skipped += 1
+                continue
+
+            (pos_tags, stems, affixes, tokens_lengths,
+             row_ids, col_ids, cell_types, ordered_cells, cell_to_positions) = enc
+
+            if gold_rc not in cell_to_positions:
+                skipped += 1
+                continue
+
+            scores, valid_cells = model(args, pos_tags, stems, affixes,
+                                        tokens_lengths, row_ids, col_ids, cell_types,
+                                        ordered_cells, cell_to_positions, device)
+            if scores is None or gold_rc not in valid_cells:
+                skipped += 1
+                continue
+
+            if atype == 'lookup':
+                pred_rc = _predict_lookup(scores, valid_cells, cells, question)
+            elif atype == 'comparison':
+                pred_rc = _predict_comparison(scores, valid_cells, cells, question)
+            else:
+                pred_rc = valid_cells[scores.argmax().item()]
+            hit     = int(pred_rc == gold_rc)
+            correct += hit
+            total   += 1
+            by_type.setdefault(atype, []).append(hit)
+
+    em = correct / total if total > 0 else 0.0
+    ts = datetime.now().strftime('%H:%M:%S')
+    print(f'  [{ts}] {split} EM={em:.4f} ({correct}/{total} correct, {skipped} skipped)')
+    for atype, hits in sorted(by_type.items()):
+        print(f'    {atype}: {sum(hits)}/{len(hits)} = {sum(hits)/len(hits):.3f}')
+    return em
+
+
+# ── Training ──────────────────────────────────────────────────────────────────
+
+def train_finetune(args, model, train_items, dev_items, kv, bpe, csv_dir, device,
+                   num_epochs=20, lr=2e-5, best_model_path=BEST_MODEL):
+    # Fine-tune top-4 seq transformer layers + cell head; freeze the rest
+    for name, p in model.named_parameters():
+        if 'cell_head' in name:
+            p.requires_grad = True
+        elif 'seq_transformer_encoder.layers.' in name:
+            # Extract layer index
+            try:
+                layer_idx = int(name.split('seq_transformer_encoder.layers.')[1].split('.')[0])
+                total_layers = model.encoder.seq_tr_nlayers \
+                    if hasattr(model.encoder, 'seq_tr_nlayers') else 12
+                p.requires_grad = (layer_idx >= total_layers - 4)
+            except (IndexError, ValueError):
+                p.requires_grad = False
+        else:
+            p.requires_grad = False
+
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(f'[FT] Trainable parameters: {trainable:,}')
+
+    optimizer = torch.optim.AdamW(
+        [p for p in model.parameters() if p.requires_grad],
+        lr=lr, weight_decay=0.01,
+    )
+    criterion = nn.CrossEntropyLoss()
+
+    best_em  = 0.0
+    history  = []
+
+    for epoch in range(1, num_epochs + 1):
+        model.train()
+        random.shuffle(train_items)
+        total_loss = 0.0
+        n_ok       = 0
+        n_skip     = 0
+
+        for item in train_items:
+            csv_path = os.path.join(csv_dir, item['table_file'])
+            if not os.path.exists(csv_path):
+                n_skip += 1
+                continue
+            cells = serialize_csv(csv_path)
+            if not cells:
+                n_skip += 1
+                continue
+
+            # Find gold cell by text matching (bypasses broken coord system)
+            gold_rc = find_gold_cell(cells, item['answer'])
+            if gold_rc is None:
+                n_skip += 1
+                continue
+
+            enc = encode_item(cells, item['question'], kv, bpe)
+            if enc is None:
+                n_skip += 1
+                continue
+
+            (pos_tags, stems, affixes, tokens_lengths,
+             row_ids, col_ids, cell_types, ordered_cells, cell_to_positions) = enc
+
+            if gold_rc not in cell_to_positions:
+                n_skip += 1
+                continue
+
+            scores, valid_cells = model(args, pos_tags, stems, affixes,
+                                        tokens_lengths, row_ids, col_ids, cell_types,
+                                        ordered_cells, cell_to_positions, device)
+            if scores is None or gold_rc not in valid_cells:
+                n_skip += 1
+                continue
+
+            gold_idx = torch.tensor([valid_cells.index(gold_rc)],
+                                    dtype=torch.long, device=device)
+            loss = criterion(scores.unsqueeze(0), gold_idx)
+
+            optimizer.zero_grad()
+            loss.backward()
+            clip_grad_norm_(model.parameters(), 1.0)
+            optimizer.step()
+
+            total_loss += loss.item()
+            n_ok       += 1
+
+        ts = datetime.now().strftime('%H:%M:%S')
+        avg_loss = total_loss / max(n_ok, 1)
+        print(f'[{ts}] Epoch {epoch}/{num_epochs}  loss={avg_loss:.4f}  '
+              f'trained={n_ok}  skipped={n_skip}')
+
+        em = evaluate(model, args, dev_items, kv, bpe, csv_dir, device, split='dev')
+        history.append({'epoch': epoch, 'train_loss': avg_loss, 'dev_em': em})
+
+        if em > best_em:
+            best_em = em
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model.state_dict(),
+                'dev_em': em,
+            }, best_model_path)
+            print(f'  ** New best EM={best_em:.4f}  saved to {best_model_path}')
+
+    return best_em, history
+
+
+# ── Main ──────────────────────────────────────────────────────────────────────
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--checkpoint', type=str, default=CHECKPOINT,
+                        help='Path to pre-trained TabuLM checkpoint (.pt)')
+    parser.add_argument('--output-prefix', type=str, default='finetune_tabqa',
+                        help='Prefix for best-model and results files (no extension)')
+    cli = parser.parse_args()
+
+    best_model_path = os.path.join(DATA_DIR, f'{cli.output_prefix}_best.pt')
+    results_path    = os.path.join(DATA_DIR, f'{cli.output_prefix}_results.json')
+
+    random.seed(42)
+    np.random.seed(42)
+    torch.manual_seed(42)
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f'[FT] Device: {device}')
+
+    # Load vocab + BPE
+    print('[FT] Loading vocab and BPE...')
+    bpe = yttm.BPE(model=os.path.join(DATA_DIR, 'BPE-30k.mdl'))
+    kv  = KBVocab()
+    kv.load_state_dict(torch.load(os.path.join(DATA_DIR, 'kb_vocab_state_dict_2021-02-07.pt'),
+                                  map_location='cpu'))
+
+    # Load checkpoint and detect architecture
+    print(f'[FT] Loading checkpoint from {cli.checkpoint}')
+    ckpt  = torch.load(cli.checkpoint, map_location='cpu')
+    state = ckpt['model_state_dict']
+    if all(k.startswith('module.') for k in state):
+        state = {k[len('module.'):]: v for k, v in state.items()}
+
+    args = detect_arch_from_state(state)
+    print(f'[FT] Detected: d_model={args.seq_tr_d_model}  nlayers={args.seq_tr_nlayers}  '
+          f'nhead={args.seq_tr_nhead}  ff={args.seq_tr_dim_feedforward}  '
+          f'morpho_dim={args.morpho_dim}  stem_dim={args.stem_dim}  '
+          f'bow={args.use_affix_bow_m_embedding}')
+
+    # Build model and load weights
+    tabulm = tabulm_base(kv, None, None, device, args, saved_model_file=None)
+    missing, unexpected = tabulm.load_state_dict(state, strict=False)
+    print(f'[FT] Loaded encoder: {len(missing)} missing, {len(unexpected)} unexpected keys')
+
+    model = TabQAModel(tabulm).to(device)
+
+    # Load and split TabQA-kin
+    print(f'[FT] Loading TabQA-kin from {TABQA_FILE}')
+    with open(TABQA_FILE) as f:
+        all_items = json.load(f)
+    print(f'[FT] {len(all_items)} QA items total')
+
+    random.shuffle(all_items)
+    split_n     = int(0.8 * len(all_items))
+    train_items = all_items[:split_n]
+    dev_items   = all_items[split_n:]
+    print(f'[FT] Train={len(train_items)}  Dev={len(dev_items)}')
+
+    print('[FT] Baseline (pre-training weights, no task training):')
+    evaluate(model, args, dev_items, kv, bpe, CSV_DIR, device, split='dev-baseline')
+
+    best_em, history = train_finetune(
+        args, model, train_items, dev_items, kv, bpe, CSV_DIR, device,
+        num_epochs=20, lr=2e-5, best_model_path=best_model_path,
+    )
+
+    print(f'\n[FT] Done. Best dev EM = {best_em:.4f}')
+
+    with open(results_path, 'w') as f:
+        json.dump({'best_em': best_em, 'epochs': history,
+                   'train_size': len(train_items), 'dev_size': len(dev_items)}, f, indent=2)
+    print(f'[FT] Results saved to {results_path}')
+
+
+if __name__ == '__main__':
+    main()

code/kinyabert_utils.py

@@ -0,0 +1,23 @@
+# Copyright (c) Antoine Nzeyimana.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from datetime import datetime
+
+def read_lines(file_name):
+    f = open(file_name, 'r')
+    lines = [line.rstrip('\n') for line in f]
+    if len(lines[-1]) == 0:
+        lines = lines[:-1]
+    if len(lines[-1]) == 0:
+        lines = lines[:-1]
+    if len(lines[-1]) == 0:
+        lines = lines[:-1]
+    if len(lines[-1]) == 0:
+        lines = lines[:-1]
+    f.close()
+    return lines
+
+def time_now():
+    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")

code/lamb.py

@@ -0,0 +1,148 @@
+# MIT License
+#
+# Copyright (c) 2019 cybertronai
+#
+# 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.
+
+"""Lamb optimizer."""
+# From: https://github.com/cybertronai/pytorch-lamb
+
+import collections
+import math
+
+import torch
+from tensorboardX import SummaryWriter
+from torch.optim import Optimizer
+
+
+def log_lamb_rs(optimizer: Optimizer, event_writer: SummaryWriter, token_count: int):
+    """Log a histogram of trust ratio scalars in across layers."""
+    results = collections.defaultdict(list)
+    for group in optimizer.param_groups:
+        for p in group['params']:
+            state = optimizer.state[p]
+            for i in ('weight_norm', 'adam_norm', 'trust_ratio'):
+                if i in state:
+                    results[i].append(state[i])
+
+    for k, v in results.items():
+        event_writer.add_histogram(f'lamb/{k}', torch.tensor(v), token_count)
+
+class Lamb(Optimizer):
+    r"""Implements Lamb algorithm.
+
+    It has been proposed in `Large Batch Optimization for Deep Learning: Training BERT in 76 minutes`_.
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        adam (bool, optional): always use trust ratio = 1, which turns this into
+            Adam. Useful for comparison purposes.
+
+    .. _Large Batch Optimization for Deep Learning: Training BERT in 76 minutes:
+        https://arxiv.org/abs/1904.00962
+    """
+
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-6,
+                 weight_decay=0, adam=False):
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+        defaults = dict(lr=lr, betas=betas, eps=eps,
+                        weight_decay=weight_decay)
+        self.adam = adam
+        super(Lamb, self).__init__(params, defaults)
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError('Lamb does not support sparse gradients, consider SparseAdam instad.')
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state['step'] = 0
+                    # Exponential moving average of gradient values
+                    state['exp_avg'] = torch.zeros_like(p.data)
+                    # Exponential moving average of squared gradient values
+                    state['exp_avg_sq'] = torch.zeros_like(p.data)
+
+                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+                beta1, beta2 = group['betas']
+
+                state['step'] += 1
+
+                # Decay the first and second moment running average coefficient
+                # m_t
+                exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
+                # v_t
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+
+                # Paper v3 does not use debiasing.
+                # bias_correction1 = 1 - beta1 ** state['step']
+                # bias_correction2 = 1 - beta2 ** state['step']
+                # Apply bias to lr to avoid broadcast.
+                step_size = group['lr'] # * math.sqrt(bias_correction2) / bias_correction1
+
+                weight_norm = p.data.pow(2).sum().sqrt().clamp(0, 10)
+
+                adam_step = exp_avg / exp_avg_sq.sqrt().add(group['eps'])
+                if group['weight_decay'] != 0:
+                    adam_step.add_(p.data, alpha=group['weight_decay'])
+
+                adam_norm = adam_step.pow(2).sum().sqrt()
+                if weight_norm == 0 or adam_norm == 0:
+                    trust_ratio = 1
+                else:
+                    trust_ratio = weight_norm / adam_norm
+                state['weight_norm'] = weight_norm
+                state['adam_norm'] = adam_norm
+                state['trust_ratio'] = trust_ratio
+                if self.adam:
+                    trust_ratio = 1
+
+                p.data.add_(adam_step, alpha=-step_size * trust_ratio)
+
+        return loss

code/morpho_common.py

@@ -0,0 +1,391 @@
+# Copyright (c) Antoine Nzeyimana.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+from cffi import FFI
+import youtokentome as yttm
+
+class SimpleKBVocab:
+
+    def __init__(self):
+        self.pos_tag_vocab = dict()
+        self.pos_tag_vocab_idx = dict()
+        self.pos_tag_vocab_idx_counts = dict()
+        self.pos_tag_vocab_idx_subsample_weights = dict()
+        self.pos_tag_vocab_idx_subsample_weights_max = 0.1
+        self.pos_tag_vocab_idx_subsample_weights_min = 0.1
+
+        self._stem_vocab = dict()
+        self._stem_vocab_idx = dict()
+        self._stem_vocab_idx_counts = dict()
+
+        self.reduced_stem_vocab = dict()
+        self.reduced_stem_vocab_idx = dict()
+        self.mapped_stem_vocab_idx = dict()
+        self.reduced_stem_vocab_idx_counts = dict()
+        self.reduced_stem_vocab_idx_subsample_weights = dict()
+        self.reduced_stem_vocab_idx_subsample_weights_max = 0.1
+        self.reduced_stem_vocab_idx_subsample_weights_min = 0.1
+
+        self.morpheme_slot_vocab = dict()
+        self.morpheme_slot_vocab_idx = dict()
+        self.morpheme_slot_vocab_idx_counts = dict()
+        self.morpheme_slot_vocab_idx_subsample_weights = dict()
+        self.morpheme_slot_vocab_idx_subsample_weights_max = 0.1
+        self.morpheme_slot_vocab_idx_subsample_weights_min = 0.1
+
+        self.affix_vocab = dict()
+        self.affix_vocab_idx = dict()
+        self.affix_vocab_idx_counts = dict()
+        self.affix_vocab_idx_subsample_weights = dict()
+        self.affix_vocab_idx_subsample_weights_max = 0.1
+        self.affix_vocab_idx_subsample_weights_min = 0.1
+
+    def load_state_dict(self, d):
+        self.pos_tag_vocab = d['pos_tag_vocab']
+        self.pos_tag_vocab_idx = d['pos_tag_vocab_idx']
+        self.pos_tag_vocab_idx_counts = d['pos_tag_vocab_idx_counts']
+        self._stem_vocab = d['_stem_vocab']
+        self._stem_vocab_idx = d['_stem_vocab_idx']
+        self._stem_vocab_idx_counts = d['_stem_vocab_idx_counts']
+        self.reduced_stem_vocab = d['reduced_stem_vocab']
+        self.mapped_stem_vocab_idx = d['mapped_stem_vocab_idx']
+        self.reduced_stem_vocab_idx_counts = d['reduced_stem_vocab_idx_counts']
+        self.morpheme_slot_vocab = d['morpheme_slot_vocab']
+        self.morpheme_slot_vocab_idx = d['morpheme_slot_vocab_idx']
+        self.morpheme_slot_vocab_idx_counts = d['morpheme_slot_vocab_idx_counts']
+        self.affix_vocab = d['affix_vocab']
+        self.affix_vocab_idx = d['affix_vocab_idx']
+        self.affix_vocab_idx_counts = d['affix_vocab_idx_counts']
+
+        self.morpheme_slot_vocab_idx_subsample_weights = d['morpheme_slot_vocab_idx_subsample_weights']
+        self.morpheme_slot_vocab_idx_subsample_weights_max = d['morpheme_slot_vocab_idx_subsample_weights_max']
+        self.morpheme_slot_vocab_idx_subsample_weights_min = d['morpheme_slot_vocab_idx_subsample_weights_min']
+
+        self.pos_tag_vocab_idx_subsample_weights = d['pos_tag_vocab_idx_subsample_weights']
+        self.pos_tag_vocab_idx_subsample_weights_max = d['pos_tag_vocab_idx_subsample_weights_max']
+        self.pos_tag_vocab_idx_subsample_weights_min = d['pos_tag_vocab_idx_subsample_weights_min']
+
+        self.affix_vocab_idx_subsample_weights = d['affix_vocab_idx_subsample_weights']
+        self.affix_vocab_idx_subsample_weights_max = d['affix_vocab_idx_subsample_weights_max']
+        self.affix_vocab_idx_subsample_weights_min = d['affix_vocab_idx_subsample_weights_min']
+
+        self.reduced_stem_vocab_idx_subsample_weights = d['reduced_stem_vocab_idx_subsample_weights']
+        self.reduced_stem_vocab_idx_subsample_weights_max = d['reduced_stem_vocab_idx_subsample_weights_max']
+        self.reduced_stem_vocab_idx_subsample_weights_min = d['reduced_stem_vocab_idx_subsample_weights_min']
+
+        for k in self.reduced_stem_vocab:
+            self.reduced_stem_vocab_idx[self.reduced_stem_vocab[k]] = k
+        self.reduced_stem_vocab_idx_counts = dict()
+        for i in self._stem_vocab_idx_counts:
+            self.reduced_stem_vocab_idx_counts[self.mapped_stem_vocab_idx[i]] = self._stem_vocab_idx_counts[i]
+
+class RichParsedToken:
+    def __init__(self, surface_form, raw_surface_form, is_apostrophed, parsed_token=None, decode_prob=None, tf_idf=0.0, pos_tag_id=None, stem_ids=None, line_num=0):
+        self.surface_form = surface_form
+        self.raw_surface_form = raw_surface_form
+        self.is_apostrophed = is_apostrophed
+        self.tf_idf = tf_idf
+        self.morpho_tokens = []
+        if parsed_token is not None:
+            parts = parsed_token.split('/')
+            self.decode_prob = float(parts[0])
+            self.tf_idf = float(parts[1])
+            morphs = parts[2].split(',')
+            pos_stem = morphs[0].split(':')
+            stem_parts = pos_stem[1].split('*')
+            if(len(stem_parts[0]) < 1):
+                # print('\nParsing wrong token: /{}/ at line # {}'.format(parsed_token, line_num))
+                self.pos_tag_idx = int(pos_stem[0])
+                self.stem_idx = [6]
+                self.morpho_slots_idx = []
+                self.affixes_idx = []
+            else:
+                self.pos_tag_idx = int(pos_stem[0])
+                self.stem_idx = [int(v) for v in stem_parts]
+                self.morpho_slots_idx = [int(morphs[i].split(':')[0]) for i in range(1, len(morphs))]
+                self.affixes_idx = [int(morphs[i].split(':')[1]) for i in range(1, len(morphs))]
+        else:
+            self.decode_prob = decode_prob
+            self.pos_tag_idx = pos_tag_id
+            self.stem_idx = stem_ids
+            self.morpho_slots_idx = []
+            self.affixes_idx = []
+
+    def append_morpheme(self, morpho_slot_id, affix_id):
+        self.morpho_slots_idx.append(morpho_slot_id)
+        self.affixes_idx.append(affix_id)
+
+    def to_parsed_format(self):
+        st = ['{}:{}'.format(self.pos_tag_idx,'*'.join([str(i) for i in self.stem_idx]))]
+        for i in range(len(self.morpho_slots_idx)):
+            st.append('{}:{}'.format(self.morpho_slots_idx[i], self.affixes_idx[i]))
+        return '{:.3g}/{:.3g}/{}'.format(self.decode_prob, self.tf_idf, ','.join(st))
+
+    def affix_set_key(self):
+        key = '-'.join([str(af) for af in self.affixes_idx]) if (len(self.affixes_idx) > 0) else 'N/A'
+        return key
+
+def build_kinlp_morpho_lib():
+    ffibuilder = FFI()
+
+    ffibuilder.cdef("""
+
+    typedef struct _snt_morpheme {
+        int slot_id;
+        int morph_id;
+        int morph_token_len;
+        char * morph_token;
+    } snt_morpheme_t;
+
+    typedef struct _snt_word {
+        char * pos_tag;
+        char * pos_group;
+        char * surface_form;
+        char * raw_surface_form;
+        char * stem;
+        snt_morpheme_t * morphemes;
+        double decode_prob;
+        double tf_idf;
+        int pos_tag_id;
+        int word_type;
+        int stem_start_index;
+        int stem_end_index;
+        int stem_start_slot_id;
+        int stem_end_slot_id;
+        int morphemes_len;
+        int apostrophed;
+    } snt_word_t;
+
+    typedef struct _snt_sentence {
+        snt_word_t * words;
+        int words_len;
+    } snt_sentence_t;
+
+    void release_sentence(snt_sentence_t * sentences, int num_sent[1]);
+    snt_sentence_t * parse_sentences_batch(const char * text, int num_sent[1]);
+
+    void start_kinlp_lib(const char * config_file);
+    void stop_kinlp_lib(void);
+
+    int is_word_morphological(const char * word);
+
+    """)
+
+    ffibuilder.set_source("kinlpmorpholib",
+                          """
+                               #include "/home/user/projects/user/kinlp/kinlp/lib.h"
+                               #include "/home/user/projects/user/kinlp/kinlp/snt.h"
+                          """,
+                          extra_compile_args=['-fopenmp', '-D use_openmp', '-O3', '-march=native', '-ffast-math'],
+                          extra_link_args=['-fopenmp'],
+                          libraries=['kinlp'])  # library name, for the linker
+
+    ffibuilder.compile(verbose=True)
+
+def parse_raw_text_lines(doc_lines, kb_vocab, bpe):
+    from kinlpmorpholib import ffi, lib
+    parsed_tokens = []
+    grouped_parsed_tokens = []
+    num_sent = ffi.new("int[1]")
+    sentences = lib.parse_sentences_batch(' '.join(doc_lines).encode('utf-8'), num_sent)
+    for i in range(num_sent[0]):
+        sent = sentences[i]
+        sentence_parsed_tokens = []
+        for j in range(sent.words_len):
+            w = sent.words[j]
+            POS_TAG = ffi.string(w.pos_tag).decode("utf-8")
+            WORD_TYPE = ffi.string(w.pos_group).decode("utf-8")
+            STEM = ffi.string(w.stem).decode("utf-8")
+            SURFACE_FORM = ffi.string(w.surface_form).decode("utf-8")
+            RAW_SURFACE_FORM = ffi.string(w.raw_surface_form).decode("utf-8")
+            DECODE_PROB = w.decode_prob
+            IS_APOSTROPHED = not (w.apostrophed == 0)
+            if IS_APOSTROPHED and (RAW_SURFACE_FORM[-1] == 'a'):
+                RAW_SURFACE_FORM = RAW_SURFACE_FORM[:-1]+"'"
+            RAW_SURFACE_FORM = RAW_SURFACE_FORM.replace('“','"').replace('‘‘', '"').replace('’’', '"').replace('’', '\'').replace('‘','\'')
+            TF_IDF = w.tf_idf
+            pti = kb_vocab.pos_tag_vocab['<UNK>']
+            if POS_TAG in kb_vocab.pos_tag_vocab.keys():
+               pti = kb_vocab.pos_tag_vocab[POS_TAG]
+
+            sids = []
+            morpho_tokens = []
+            stem_tokens = []
+            if ((STEM == SURFACE_FORM) and (w.morphemes_len <= 0)):
+                list_sub_words = bpe.encode(SURFACE_FORM, output_type=yttm.OutputType.SUBWORD)
+                for sub_word in list_sub_words:
+                    stem_key = WORD_TYPE + ":" + sub_word
+                    morpho_tokens.append(stem_key)
+                    stem_tokens.append(stem_key)
+                    si = kb_vocab._stem_vocab['<UNK>']
+                    if stem_key in kb_vocab._stem_vocab.keys():
+                        si = kb_vocab._stem_vocab[stem_key]
+                    sids.append(si)
+            else:
+                stem_key = WORD_TYPE + ":" + STEM
+                morpho_tokens.append(stem_key)
+                stem_tokens.append(stem_key)
+                si = kb_vocab._stem_vocab['<UNK>']
+                if stem_key in kb_vocab._stem_vocab.keys():
+                    si = kb_vocab._stem_vocab[stem_key]
+                sids.append(si)
+            ptoken = RichParsedToken(SURFACE_FORM, RAW_SURFACE_FORM, IS_APOSTROPHED, parsed_token=None, decode_prob=DECODE_PROB, tf_idf=TF_IDF, pos_tag_id=pti, stem_ids=sids)
+            if (w.morphemes_len > 0):
+                for k in range(w.morphemes_len):
+                    if ((k != w.stem_start_index) and (k != w.stem_end_index)):
+                        MORPHEME_SLOT = WORD_TYPE + ":" + str(w.morphemes[k].slot_id)
+                        MORPHEME = MORPHEME_SLOT + ":" + ffi.string(w.morphemes[k].morph_token).decode("utf-8")
+
+                        if (k < w.stem_start_index):
+                            morpho_tokens = [MORPHEME] + morpho_tokens
+                        elif (k < w.stem_end_index) and (len(stem_tokens) == 1) and (w.stem_start_index != w.stem_end_index):
+                            # For reduplication for verbs, a bit rare
+                            morpho_tokens = morpho_tokens + [MORPHEME] + stem_tokens
+                        else:
+                            morpho_tokens = morpho_tokens + [MORPHEME]
+
+                        msi = kb_vocab.morpheme_slot_vocab['<UNK>']
+                        if MORPHEME_SLOT in kb_vocab.morpheme_slot_vocab.keys():
+                            msi = kb_vocab.morpheme_slot_vocab[MORPHEME_SLOT]
+
+                        mi = kb_vocab.affix_vocab['<UNK>']
+                        if MORPHEME in kb_vocab.affix_vocab.keys():
+                            mi = kb_vocab.affix_vocab[MORPHEME]
+
+                        ptoken.append_morpheme(msi, mi)
+            ptoken.morpho_tokens.extend(morpho_tokens)
+            parsed_tokens.append(ptoken)
+            sentence_parsed_tokens.append(ptoken)
+        grouped_parsed_tokens.append(sentence_parsed_tokens)
+    lib.release_sentence(sentences, num_sent)
+
+    return parsed_tokens, grouped_parsed_tokens
+
+def str2bool(v):
+    if isinstance(v, bool):
+        return v
+    if v.lower() in ('yes', 'true', 't', 'y', '1'):
+        return True
+    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
+        return False
+    else:
+        raise argparse.ArgumentTypeError('Boolean value expected.')
+
+
+def setup_common_args(list_args=None, silent=False):
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-g', '--gpus', default=1, type=int, help='number of gpus')
+    parser.add_argument('-p', '--pos', default=1, type=int, help='number of POS embeddings for morphology')
+    parser.add_argument('-s', '--stem', default=1, type=int, help='number of Stem embeddings for morphology')
+    parser.add_argument('-ii', '--inference-iters', default=1, type=int, help='number of MLM Inference iterations')
+    parser.add_argument('-ir', '--inference-runs', default=20, type=int,
+                        help='number of MLM Inference runes per iteration')
+    parser.add_argument("--use-affix-bow", type=str2bool, default=False,
+                        help="Use affix embeddings sum (BOW) for morphology")
+    parser.add_argument("--use-pos-aware-rel", type=str2bool, default=True,
+                        help="Use POS-aware relative position embedding")
+    parser.add_argument("--use-tupe-rel", type=str2bool, default=False, help="Use TUPE relative position embedding")
+    parser.add_argument("--resume-from-best-saved", type=str2bool, default=False, help="Resume training from best saved model")
+
+    parser.add_argument("--use-afsets", type=str2bool, default=False)
+    parser.add_argument("--predict-affixes", type=str2bool, default=True)
+
+    # KinyaBERT_base architecture hyper-parameters
+    parser.add_argument("--seq-tr-dropout", type=float, default=0.1)
+    parser.add_argument("--layernorm-epsilon", type=float, default=1e-6)
+    parser.add_argument("--morpho-tr-dropout", type=float, default=0.1)
+    parser.add_argument("--morpho-tr-nhead", type=int, default=4)
+    parser.add_argument("--morpho-tr-nlayers", type=int, default=4)
+    parser.add_argument("--morpho-dim", type=int, default=128)
+    parser.add_argument("--morpho-tr-dim-feedforward", type=int, default=512)
+    parser.add_argument("--stem-dim", type=int, default=256)
+    parser.add_argument("--max-seq-len", type=int, default=512)
+    parser.add_argument("--seq-tr-nhead", type=int, default=12)
+    parser.add_argument("--seq-tr-nlayers", type=int, default=12)
+    parser.add_argument("--seq-tr-dim-feedforward", type=int, default=3072)
+
+    parser.add_argument("--batch-size", type=int, default=20)
+    parser.add_argument("--num-epochs", type=int, default=20)
+    parser.add_argument("--accumulation-steps", type=int, default=128)
+    parser.add_argument("--number-of-load-batches", type=int, default=384)
+    parser.add_argument("--max-input-lines", type=int, default=999999)
+    parser.add_argument("--warmup-ratio", type=float, default=0.06)
+
+    parser.add_argument("--num-iters", type=int, default=200000)
+    parser.add_argument("--warmup-iter", type=int, default=2000)
+
+    parser.add_argument("--peak-lr", type=float, default=4e-4)
+    parser.add_argument("--wd", type=float, default=0.01)
+
+    parser.add_argument("--cls-labels", type=str, default="0,1")
+
+    parser.add_argument("--cls-train-input0", type=str, default=None)
+    parser.add_argument("--cls-train-input1", type=str, default=None)
+    parser.add_argument("--cls-train-label", type=str, default=None)
+
+    parser.add_argument("--cls-dev-input0", type=str, default=None)
+    parser.add_argument("--cls-dev-input1", type=str, default=None)
+    parser.add_argument("--cls-dev-label", type=str, default=None)
+
+    parser.add_argument("--cls-test-input0", type=str, default=None)
+    parser.add_argument("--cls-test-input1", type=str, default=None)
+    parser.add_argument("--cls-test-label", type=str, default=None)
+
+    parser.add_argument("--pretrained-model-file", type=str, default=None)
+    parser.add_argument("--devbest-cls-model-save-file-path", type=str, default=None)
+    parser.add_argument("--final-cls-model-save-file-path", type=str, default=None)
+
+    parser.add_argument("--devbest-cls-output-file", type=str, default=None)
+    parser.add_argument("--final-cls-output-file", type=str, default=None)
+
+    parser.add_argument("--home-path", type=str, default="./")
+
+    parser.add_argument("--regression-target", type=str2bool, default=False)
+    parser.add_argument("--regression-scale-factor", type=float, default=5.0)
+
+    parser.add_argument("--pretrained-roberta-model-dir", type=str, default="data/")
+    parser.add_argument("--pretrained-roberta-checkpoint-file", type=str, default="checkpoint_best.pt")
+    parser.add_argument("--xlmr", type=str2bool, default=False)
+
+    parser.add_argument("--pooler-dropout", type=float, default=0.1)
+
+    parser.add_argument("--embed-dim", type=int, default=768)
+
+    parser.add_argument("--inference-model-file", type=str, default=None)
+
+    parser.add_argument("--model-keyword", type=str, default=None)
+    parser.add_argument("--task-keyword", type=str, default=None)
+    parser.add_argument("--input-format", type=str, default=None)
+
+    parser.add_argument("--afset-dict-size", type=int, default=10000)
+
+    parser.add_argument("--debug", type=str2bool, default=False)
+
+    parser.add_argument("--use-morpho-encoder", type=str2bool, default=True)
+
+    parser.add_argument("--exploratory-model-load", type=str, default=None)
+
+    parser.add_argument("-f", type=str, default=None)
+
+
+    if list_args is not None:
+        args = parser.parse_args(list_args)
+    else:
+        args = parser.parse_args()
+
+    args.world_size = args.gpus
+    args.num_pos_m_embeddings = args.pos
+    args.num_stem_m_embeddings = args.stem
+    args.use_affix_bow_m_embedding = args.use_affix_bow
+    args.use_pos_aware_rel_pos_bias = args.use_pos_aware_rel
+    args.use_tupe_rel_pos_bias = args.use_tupe_rel
+    args.num_inference_iters = args.inference_iters
+    args.num_inference_runs = args.inference_runs
+
+    if not silent:
+        print('Call arguments:\n', args)
+
+    return args

code/morpho_data_loaders.py

@@ -0,0 +1,958 @@
+# Copyright (c) Antoine Nzeyimana.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+from __future__ import print_function, division
+
+import math
+import random
+import sys
+# Ignore warnings
+import warnings
+from typing import List
+
+import numpy as np
+import progressbar
+import torch
+import youtokentome as yttm
+from torch.utils.data import Dataset
+
+from kinyabert_utils import time_now
+
+warnings.filterwarnings("ignore")
+
+def read_corpus(fn):
+    f = open(fn, 'r+')
+    corpus_lines = [line.rstrip('\n') for line in f]
+    f.close()
+    corpus_lines.append("\n")
+    return corpus_lines
+
+def read_vocab_idx(fn, voc, voc_idx):
+    f = open(fn, 'r')
+    v_lines = f.readlines()
+    f.close()
+    for v in v_lines:
+        if(len(v)>1):
+            vt = v.split('\t')
+            if(len(vt) == 2):
+                voc_idx[int(vt[1])] = vt[0]
+                voc[vt[0]] = int(vt[1])
+    return voc, voc_idx
+
+def read_vocab_counts(fn, voc, voc_idx_counts):
+    f = open(fn, 'r')
+    v_lines = f.readlines()
+    f.close()
+    for v in v_lines:
+        if(len(v)>1):
+            vt = v.split('\t')
+            if(len(vt) == 2):
+                id = voc[vt[0]]
+                voc_idx_counts[id] = int(vt[1])
+    return voc, voc_idx_counts
+
+def reduce_stem_vocab(_stem_vocab, _stem_vocab_idx, _stem_vocab_idx_counts,
+                      noun_min_count, verb_min_count, np_min_count,
+                      other_morpho_min_count, other_cls_min_count, other_token_min_count):
+    reduced_stem_vocab = dict()
+    mapped_stem_vocab_idx = dict()
+    reduced_stem_vocab_idx_counts = dict()
+    unk_idx = _stem_vocab['<UNK>']
+    for i in range(1,len(_stem_vocab_idx)+1):
+        if (i <= unk_idx):
+            mapped_stem_vocab_idx[i] = i
+            reduced_stem_vocab[_stem_vocab_idx[i]] = i
+            reduced_stem_vocab_idx_counts[i] = 1000000
+        else:
+            key = _stem_vocab_idx[i]
+            if (key.startswith('N:')):
+                min_count = noun_min_count
+            elif (key.startswith('V:')):
+                min_count = verb_min_count
+            elif (key.startswith('QA:') or key.startswith('PO:') or key.startswith('DE:') or key.startswith('NU:') or key.startswith('OT:')):
+                min_count = other_morpho_min_count
+            elif (key.startswith('NP:')):
+                min_count = np_min_count
+            elif (key.startswith('T:')):
+                min_count = other_token_min_count
+            else:
+                min_count = other_cls_min_count
+            if _stem_vocab_idx_counts[i] >= min_count:
+                idx = len(reduced_stem_vocab)+1
+                mapped_stem_vocab_idx[i] = idx
+                reduced_stem_vocab[_stem_vocab_idx[i]] = idx
+                reduced_stem_vocab_idx_counts[idx] = _stem_vocab_idx_counts[i]
+            else:
+                mapped_stem_vocab_idx[i] = reduced_stem_vocab['<UNK>']
+
+    return reduced_stem_vocab, mapped_stem_vocab_idx, reduced_stem_vocab_idx_counts
+
+def sigmoid_score(x, min_val, max_val):
+    s = pow((1.0 + (math.exp(((-8.0) * (x - min_val)) / (max_val - min_val)))), -8.0)
+    return s
+
+class KBVocab:
+
+    def __init__(self, noun_min_count = 200,
+                 verb_min_count = 100,
+                 np_min_count = 200,
+                 other_morpho_min_count = 20,
+                 other_cls_min_count = 20,
+                 other_token_min_count = 200,
+                 pos_tag_vocab_idx_tsv="data/pos_tag_vocab_idx_2021-02-07.tsv",
+                 pos_tag_vocab_tsv="data/pos_tag_vocab_2021-02-07.tsv",
+                 stem_vocab_idx_tsv="data/stem_vocab_idx_2021-02-07.tsv",
+                 stem_vocab_tsv="data/stem_vocab_2021-02-07.tsv",
+                 morpheme_slot_vocab_idx_tsv="data/morpheme_slot_vocab_idx_2021-02-07.tsv",
+                 morpheme_slot_vocab_tsv="data/morpheme_slot_vocab_2021-02-07.tsv",
+                 affix_vocab_idx_tsv="data/affix_vocab_idx_2021-02-07.tsv",
+                 affix_vocab_tsv="data/affix_vocab_2021-02-07.tsv",
+                 read_vocab_files = False):
+        self.pos_tag_vocab = dict()
+        self.pos_tag_vocab_idx = dict()
+        self.pos_tag_vocab_idx_counts = dict()
+        self.pos_tag_vocab_idx_subsample_weights = dict()
+        self.pos_tag_vocab_idx_subsample_weights_max = 0.1
+        self.pos_tag_vocab_idx_subsample_weights_min = 0.1
+        if(read_vocab_files):
+            read_vocab_idx(pos_tag_vocab_idx_tsv, self.pos_tag_vocab, self.pos_tag_vocab_idx)
+            read_vocab_counts(pos_tag_vocab_tsv, self.pos_tag_vocab, self.pos_tag_vocab_idx_counts)
+
+        self._stem_vocab = dict()
+        self._stem_vocab_idx = dict()
+        self._stem_vocab_idx_counts = dict()
+        if(read_vocab_files):
+            read_vocab_idx(stem_vocab_idx_tsv, self._stem_vocab, self._stem_vocab_idx)
+            read_vocab_counts(stem_vocab_tsv, self._stem_vocab, self._stem_vocab_idx_counts)
+
+        self.reduced_stem_vocab = dict()
+        self.reduced_stem_vocab_idx = dict()
+        self.mapped_stem_vocab_idx = dict()
+        self.reduced_stem_vocab_idx_counts = dict()
+        self.reduced_stem_vocab_idx_subsample_weights = dict()
+        self.reduced_stem_vocab_idx_subsample_weights_max = 0.1
+        self.reduced_stem_vocab_idx_subsample_weights_min = 0.1
+        if(read_vocab_files):
+            (self.reduced_stem_vocab,
+             self.mapped_stem_vocab_idx,
+             self.reduced_stem_vocab_idx_counts) = reduce_stem_vocab(self._stem_vocab,
+                                                                     self._stem_vocab_idx,
+                                                                     self._stem_vocab_idx_counts,
+                                                                     noun_min_count,
+                                                                     verb_min_count,
+                                                                     np_min_count,
+                                                                     other_morpho_min_count,
+                                                                     other_cls_min_count,
+                                                                     other_token_min_count)
+            for k in self.reduced_stem_vocab:
+                self.reduced_stem_vocab_idx[self.reduced_stem_vocab[k]] = k
+
+        self.morpheme_slot_vocab = dict()
+        self.morpheme_slot_vocab_idx = dict()
+        self.morpheme_slot_vocab_idx_counts = dict()
+        self.morpheme_slot_vocab_idx_subsample_weights = dict()
+        self.morpheme_slot_vocab_idx_subsample_weights_max = 0.1
+        self.morpheme_slot_vocab_idx_subsample_weights_min = 0.1
+        if(read_vocab_files):
+            read_vocab_idx(morpheme_slot_vocab_idx_tsv, self.morpheme_slot_vocab, self.morpheme_slot_vocab_idx)
+            read_vocab_counts(morpheme_slot_vocab_tsv, self.morpheme_slot_vocab, self.morpheme_slot_vocab_idx_counts)
+
+        self.affix_vocab = dict()
+        self.affix_vocab_idx = dict()
+        self.affix_vocab_idx_counts = dict()
+        self.affix_vocab_idx_subsample_weights = dict()
+        self.affix_vocab_idx_subsample_weights_max = 0.1
+        self.affix_vocab_idx_subsample_weights_min = 0.1
+        if(read_vocab_files):
+            read_vocab_idx(affix_vocab_idx_tsv, self.affix_vocab, self.affix_vocab_idx)
+            read_vocab_counts(affix_vocab_tsv, self.affix_vocab, self.affix_vocab_idx_counts)
+
+    def compute_subsampling_weights(self):
+        eps = 1e-5
+        sum_pos_tags = sum([self.pos_tag_vocab_idx_counts[k] for k in self.pos_tag_vocab_idx_counts])
+        for k in self.pos_tag_vocab_idx_counts:
+            z = float(self.pos_tag_vocab_idx_counts[k])/float(sum_pos_tags)
+            z = z + eps
+            self.pos_tag_vocab_idx_subsample_weights[k] = (math.sqrt(z/0.001)+1.0)*(0.001/z)
+        self.pos_tag_vocab_idx_subsample_weights_max = max([self.pos_tag_vocab_idx_subsample_weights[k] for k in self.pos_tag_vocab_idx_subsample_weights])
+        self.pos_tag_vocab_idx_subsample_weights_min = min([self.pos_tag_vocab_idx_subsample_weights[k] for k in self.pos_tag_vocab_idx_subsample_weights])
+        for k in self.pos_tag_vocab_idx_subsample_weights:
+            if (k <= self.pos_tag_vocab['<UNK>']):
+                self.pos_tag_vocab_idx_subsample_weights[k] = self.pos_tag_vocab_idx_subsample_weights_min
+
+        sum_morpheme_slots = sum([self.morpheme_slot_vocab_idx_counts[k] for k in self.morpheme_slot_vocab_idx_counts])
+        for k in self.morpheme_slot_vocab_idx_counts:
+            z = float(self.morpheme_slot_vocab_idx_counts[k])/float(sum_morpheme_slots)
+            z = z + eps
+            self.morpheme_slot_vocab_idx_subsample_weights[k] = (math.sqrt(z/0.001)+1.0)*(0.001/z)
+        self.morpheme_slot_vocab_idx_subsample_weights_max = max([self.morpheme_slot_vocab_idx_subsample_weights[k] for k in self.morpheme_slot_vocab_idx_subsample_weights])
+        self.morpheme_slot_vocab_idx_subsample_weights_min = min([self.morpheme_slot_vocab_idx_subsample_weights[k] for k in self.morpheme_slot_vocab_idx_subsample_weights])
+        for k in self.morpheme_slot_vocab_idx_subsample_weights:
+            if (k <= self.morpheme_slot_vocab['<UNK>']):
+                self.morpheme_slot_vocab_idx_subsample_weights[k] = self.morpheme_slot_vocab_idx_subsample_weights_min
+
+        sum_affixes = sum([self.affix_vocab_idx_counts[k] for k in self.affix_vocab_idx_counts])
+        for k in self.affix_vocab_idx_counts:
+            z = float(self.affix_vocab_idx_counts[k])/float(sum_affixes)
+            z = z + eps
+            self.affix_vocab_idx_subsample_weights[k] = (math.sqrt(z/0.001)+1.0)*(0.001/z)
+        self.affix_vocab_idx_subsample_weights_max = max([self.affix_vocab_idx_subsample_weights[k] for k in self.affix_vocab_idx_subsample_weights])
+        self.affix_vocab_idx_subsample_weights_min = min([self.affix_vocab_idx_subsample_weights[k] for k in self.affix_vocab_idx_subsample_weights])
+        for k in self.affix_vocab_idx_subsample_weights:
+            if (k <= self.affix_vocab['<UNK>']):
+                self.affix_vocab_idx_subsample_weights[k] = self.affix_vocab_idx_subsample_weights_min
+
+        sum_stems = sum([self.reduced_stem_vocab_idx_counts[k] for k in self.reduced_stem_vocab_idx_counts])
+        for k in self.reduced_stem_vocab_idx_counts:
+            z = float(self.reduced_stem_vocab_idx_counts[k])/float(sum_stems)
+            z = z + eps
+            self.reduced_stem_vocab_idx_subsample_weights[k] = (math.sqrt(z/0.001)+1.0)*(0.001/z)
+        self.reduced_stem_vocab_idx_subsample_weights_max = max([self.reduced_stem_vocab_idx_subsample_weights[k] for k in self.reduced_stem_vocab_idx_subsample_weights])
+        self.reduced_stem_vocab_idx_subsample_weights_min = min([self.reduced_stem_vocab_idx_subsample_weights[k] for k in self.reduced_stem_vocab_idx_subsample_weights])
+        for k in self.reduced_stem_vocab_idx_subsample_weights:
+            if (k <= self.reduced_stem_vocab['<UNK>']):
+                self.reduced_stem_vocab_idx_subsample_weights[k] = self.reduced_stem_vocab_idx_subsample_weights_min
+
+    def state_dict(self):
+        return {'pos_tag_vocab':self.pos_tag_vocab,
+                'pos_tag_vocab_idx':self.pos_tag_vocab_idx,
+                'pos_tag_vocab_idx_counts':self.pos_tag_vocab_idx_counts,
+                '_stem_vocab':self._stem_vocab,
+                '_stem_vocab_idx':self._stem_vocab_idx,
+                '_stem_vocab_idx_counts':self._stem_vocab_idx_counts,
+                'reduced_stem_vocab':self.reduced_stem_vocab,
+                'mapped_stem_vocab_idx':self.mapped_stem_vocab_idx,
+                'reduced_stem_vocab_idx_counts':self.reduced_stem_vocab_idx_counts,
+                'morpheme_slot_vocab':self.morpheme_slot_vocab,
+                'morpheme_slot_vocab_idx':self.morpheme_slot_vocab_idx,
+                'morpheme_slot_vocab_idx_counts':self.morpheme_slot_vocab_idx_counts,
+                'affix_vocab':self.affix_vocab,
+                'affix_vocab_idx':self.affix_vocab_idx,
+                'affix_vocab_idx_counts':self.affix_vocab_idx_counts,
+
+                'morpheme_slot_vocab_idx_subsample_weights': self.morpheme_slot_vocab_idx_subsample_weights,
+                'morpheme_slot_vocab_idx_subsample_weights_max': self.morpheme_slot_vocab_idx_subsample_weights_max,
+                'morpheme_slot_vocab_idx_subsample_weights_min': self.morpheme_slot_vocab_idx_subsample_weights_min,
+
+                'affix_vocab_idx_subsample_weights':self.affix_vocab_idx_subsample_weights,
+                'affix_vocab_idx_subsample_weights_max':self.affix_vocab_idx_subsample_weights_max,
+                'affix_vocab_idx_subsample_weights_min':self.affix_vocab_idx_subsample_weights_min,
+
+                'pos_tag_vocab_idx_subsample_weights': self.pos_tag_vocab_idx_subsample_weights,
+                'pos_tag_vocab_idx_subsample_weights_max': self.pos_tag_vocab_idx_subsample_weights_max,
+                'pos_tag_vocab_idx_subsample_weights_min': self.pos_tag_vocab_idx_subsample_weights_min,
+
+                'reduced_stem_vocab_idx_subsample_weights':self.reduced_stem_vocab_idx_subsample_weights,
+                'reduced_stem_vocab_idx_subsample_weights_max':self.reduced_stem_vocab_idx_subsample_weights_max,
+                'reduced_stem_vocab_idx_subsample_weights_min':self.reduced_stem_vocab_idx_subsample_weights_min}
+
+    def load_state_dict(self, d):
+        self.pos_tag_vocab = d['pos_tag_vocab']
+        self.pos_tag_vocab_idx = d['pos_tag_vocab_idx']
+        self.pos_tag_vocab_idx_counts = d['pos_tag_vocab_idx_counts']
+        self._stem_vocab = d['_stem_vocab']
+        self._stem_vocab_idx = d['_stem_vocab_idx']
+        self._stem_vocab_idx_counts = d['_stem_vocab_idx_counts']
+        self.reduced_stem_vocab = d['reduced_stem_vocab']
+        self.mapped_stem_vocab_idx = d['mapped_stem_vocab_idx']
+        self.reduced_stem_vocab_idx_counts = d['reduced_stem_vocab_idx_counts']
+        self.morpheme_slot_vocab = d['morpheme_slot_vocab']
+        self.morpheme_slot_vocab_idx = d['morpheme_slot_vocab_idx']
+        self.morpheme_slot_vocab_idx_counts = d['morpheme_slot_vocab_idx_counts']
+        self.affix_vocab = d['affix_vocab']
+        self.affix_vocab_idx = d['affix_vocab_idx']
+        self.affix_vocab_idx_counts = d['affix_vocab_idx_counts']
+
+        self.morpheme_slot_vocab_idx_subsample_weights = d['morpheme_slot_vocab_idx_subsample_weights']
+        self.morpheme_slot_vocab_idx_subsample_weights_max = d['morpheme_slot_vocab_idx_subsample_weights_max']
+        self.morpheme_slot_vocab_idx_subsample_weights_min = d['morpheme_slot_vocab_idx_subsample_weights_min']
+
+        self.pos_tag_vocab_idx_subsample_weights = d['pos_tag_vocab_idx_subsample_weights']
+        self.pos_tag_vocab_idx_subsample_weights_max = d['pos_tag_vocab_idx_subsample_weights_max']
+        self.pos_tag_vocab_idx_subsample_weights_min = d['pos_tag_vocab_idx_subsample_weights_min']
+
+        self.affix_vocab_idx_subsample_weights = d['affix_vocab_idx_subsample_weights']
+        self.affix_vocab_idx_subsample_weights_max = d['affix_vocab_idx_subsample_weights_max']
+        self.affix_vocab_idx_subsample_weights_min = d['affix_vocab_idx_subsample_weights_min']
+
+        self.reduced_stem_vocab_idx_subsample_weights = d['reduced_stem_vocab_idx_subsample_weights']
+        self.reduced_stem_vocab_idx_subsample_weights_max = d['reduced_stem_vocab_idx_subsample_weights_max']
+        self.reduced_stem_vocab_idx_subsample_weights_min = d['reduced_stem_vocab_idx_subsample_weights_min']
+
+        for k in self.reduced_stem_vocab:
+            self.reduced_stem_vocab_idx[self.reduced_stem_vocab[k]] = k
+        self.reduced_stem_vocab_idx_counts = dict()
+        for i in self._stem_vocab_idx_counts:
+            self.reduced_stem_vocab_idx_counts[self.mapped_stem_vocab_idx[i]] = self._stem_vocab_idx_counts[i]
+
+class AffixSetVocab:
+    def __init__(self, reduced_affix_dict_file = None, reduced_affix_dict_map_file = None):
+        self.affix_set_vocab_idx = dict()
+        self.affix_set_vocab = dict()
+        self.reduced_affix_dict_counts = dict()
+        self.reduced_affix_dict_map = dict()
+
+        if reduced_affix_dict_file is not None:
+            f = open(reduced_affix_dict_file, 'r')
+            dict_lines = [line.rstrip('\n') for line in f]
+            f.close()
+            idx = 1
+            for l in dict_lines:
+                if len(l) > 0:
+                    spl = l.split(',')
+                    if (len(spl) == 2):
+                        self.reduced_affix_dict_counts[spl[0]] = int(spl[1])
+                        self.affix_set_vocab_idx[spl[0]] = idx
+                        self.affix_set_vocab[idx] = spl[0]
+                        idx += 1
+
+        if reduced_affix_dict_map_file is not None:
+            f = open(reduced_affix_dict_map_file, 'r')
+            dict_lines = [line.rstrip('\n') for line in f]
+            f.close()
+            for l in dict_lines:
+                if len(l) > 0:
+                    spl = l.split(',')
+                    if (len(spl) == 2):
+                        self.reduced_affix_dict_map[spl[0]] = spl[1]
+                        if (spl[1]) in self.affix_set_vocab_idx:
+                            self.affix_set_vocab_idx[spl[0]] = self.affix_set_vocab_idx[spl[1]]
+                        else:
+                            self.affix_set_vocab_idx[spl[0]] = 1
+
+    def affix_set_to_idx(self, key):
+        if key in self.affix_set_vocab_idx:
+            return self.affix_set_vocab_idx[key]
+        else:
+            return 1 # N/A
+
+    def affix_set_idx_to_txt(self, idx, kb_vocab: KBVocab):
+        if idx == 1:
+            return "N/A"
+        elif idx in self.affix_set_vocab:
+            return '-'.join([kb_vocab.affix_vocab_idx[int(x)] for x in self.affix_set_vocab[idx].split('-')])
+        else:
+            return "UNK" # N/A
+
+    def random_idx(self):
+        return random.randint(1, len(self.affix_set_vocab_idx))
+
+    def state_dict(self):
+        return {'affix_set_vocab_idx':self.affix_set_vocab_idx,
+                'reduced_affix_dict_counts':self.reduced_affix_dict_counts,
+                'reduced_affix_dict_map':self.reduced_affix_dict_map}
+
+    def load_state_dict(self, d):
+        self.affix_set_vocab_idx = d['affix_set_vocab_idx']
+        self.reduced_affix_dict_counts = d['reduced_affix_dict_counts']
+        self.reduced_affix_dict_map = d['reduced_affix_dict_map']
+
+class ParsedToken:
+    def __init__(self, surface_form, parsed_token=None, decode_prob=None, tf_idf=0.0, pos_tag_id=None, stem_ids=None, line_num=0):
+        self.surface_form = surface_form
+        self.tf_idf = tf_idf
+        if parsed_token is not None:
+            parts = parsed_token.split('/')
+            self.decode_prob = float(parts[0])
+            self.tf_idf = float(parts[1])
+            morphs = parts[2].split(',')
+            pos_stem = morphs[0].split(':')
+            stem_parts = pos_stem[1].split('*')
+            if(len(stem_parts[0]) < 1):
+                # print('\nParsing wrong token: /{}/ at line # {}'.format(parsed_token, line_num))
+                self.pos_tag_idx = int(pos_stem[0])
+                self.stem_idx = [6]
+                self.morpho_slots_idx = []
+                self.affixes_idx = []
+            else:
+                self.pos_tag_idx = int(pos_stem[0])
+                self.stem_idx = [int(v) for v in stem_parts]
+                self.morpho_slots_idx = [int(morphs[i].split(':')[0]) for i in range(1, len(morphs))]
+                self.affixes_idx = [int(morphs[i].split(':')[1]) for i in range(1, len(morphs))]
+        else:
+            self.decode_prob = decode_prob
+            self.pos_tag_idx = pos_tag_id
+            self.stem_idx = stem_ids
+            self.morpho_slots_idx = []
+            self.affixes_idx = []
+
+    def append_morpheme(self, morpho_slot_id, affix_id):
+        self.morpho_slots_idx.append(morpho_slot_id)
+        self.affixes_idx.append(affix_id)
+
+    def to_parsed_format(self):
+        st = ['{}:{}'.format(self.pos_tag_idx,'*'.join([str(i) for i in self.stem_idx]))]
+        for i in range(len(self.morpho_slots_idx)):
+            st.append('{}:{}'.format(self.morpho_slots_idx[i], self.affixes_idx[i]))
+        return '{:.3g}/{:.3g}/{}'.format(self.decode_prob, self.tf_idf, ','.join(st))
+
+    def affix_set_key(self):
+        key = '-'.join([str(af) for af in self.affixes_idx]) if (len(self.affixes_idx) > 0) else 'N/A'
+        return key
+
+def update_tf_idf_from_idf(doc_sentences):
+    doc_voc = dict()
+    doc_size = 0.0
+    for parsed_tokens in doc_sentences:
+        for t in parsed_tokens:
+            doc_size += 1.0
+            v = 0.0
+            if(t.stem_idx[0] in doc_voc):
+                v = doc_voc[t.stem_idx[0]]
+            doc_voc[t.stem_idx[0]] = v + 1.0
+    for parsed_tokens in doc_sentences:
+        for t in parsed_tokens:
+            t.tf_idf = sigmoid_score(t.tf_idf * doc_voc[t.stem_idx[0]] / doc_size, 0.01, 0.24)
+
+def pre_process_parsed_corpus_compute_tfidf(input_corpus, output_corpus):
+    f = open(input_corpus, 'r')
+    Lines = f.readlines()
+    f.close()
+
+    outfile = open(output_corpus, 'w')
+    doc_idx = [i for i in range(len(Lines)) if (len(Lines[i]) == 1)]
+    if doc_idx[-1] < (len(Lines) - 1):
+        doc_idx.append(len(Lines))
+    start_idx = 0
+    all_docs = len(doc_idx)
+    print_docs = 0
+
+    tot = 0
+    for end_idx in doc_idx:
+        tot = tot + 1
+        lines_batch = Lines[start_idx:end_idx]
+        start_idx = end_idx + 1
+        if (len(lines_batch) > 0):
+            doc_sentences = []
+            for ln in lines_batch:
+                line = ln.strip()
+                line = line.strip('\n')
+                line = line.strip('\t')
+                line = line.strip('\r')
+                splits = line.split('; ')
+                if (len(splits) > 0):
+                    if ((len(splits[0]) > 4) and ('/' in splits[0]) and (':' in splits[0])):
+                        doc_sentences.append([ParsedToken('_', parsed_token=t) for t in splits])
+            update_tf_idf_from_idf(doc_sentences)
+            for sent in doc_sentences:
+                outfile.write('; '.join([tok.to_parsed_format() for tok in sent]) + "\n")
+            outfile.write("\n")
+            outfile.flush()
+            print_docs += 1
+    outfile.close()
+    print('Exported: {} / {} / {}'.format(print_docs, tot, all_docs))
+
+def parse_raw_text_lines(doc_lines, kb_vocab, bpe):
+    from kinlpmorpholib import ffi, lib
+    parsed_tokens = []
+    num_sent = ffi.new("int[1]")
+    sentences = lib.parse_sentences_batch(doc_lines.encode('utf-8'), num_sent)
+    for i in range(num_sent[0]):
+        sent = sentences[i]
+        for j in range(sent.words_len):
+            w = sent.words[j]
+            POS_TAG = ffi.string(w.pos_tag).decode("utf-8")
+            WORD_TYPE = ffi.string(w.pos_group).decode("utf-8")
+            STEM = ffi.string(w.stem).decode("utf-8")
+            SURFACE_FORM = ffi.string(w.surface_form).decode("utf-8")
+            DECODE_PROB = w.decode_prob
+            TF_IDF = w.tf_idf
+            pti = kb_vocab.pos_tag_vocab['<UNK>']
+            if POS_TAG in kb_vocab.pos_tag_vocab.keys():
+               pti = kb_vocab.pos_tag_vocab[POS_TAG]
+
+            sids = []
+            if ((STEM == SURFACE_FORM) and (w.morphemes_len <= 0)):
+                list_sub_words = bpe.encode(SURFACE_FORM, output_type=yttm.OutputType.SUBWORD)
+                for sub_word in list_sub_words:
+                    stem_key = WORD_TYPE + ":" + sub_word
+                    si = kb_vocab._stem_vocab['<UNK>']
+                    if stem_key in kb_vocab._stem_vocab.keys():
+                        si = kb_vocab._stem_vocab[stem_key]
+                    sids.append(si)
+            else:
+                stem_key = WORD_TYPE + ":" + STEM
+                si = kb_vocab._stem_vocab['<UNK>']
+                if stem_key in kb_vocab._stem_vocab.keys():
+                    si = kb_vocab._stem_vocab[stem_key]
+                sids.append(si)
+            ptoken = ParsedToken(SURFACE_FORM, parsed_token=None, decode_prob=DECODE_PROB, tf_idf=TF_IDF, pos_tag_id=pti, stem_ids=sids)
+            if (w.morphemes_len > 0):
+                for k in range(w.morphemes_len):
+                    if ((k != w.stem_start_index) and (k != w.stem_end_index)):
+                        MORPHEME_SLOT = WORD_TYPE + ":" + str(w.morphemes[k].slot_id)
+                        MORPHEME = MORPHEME_SLOT + ":" + ffi.string(w.morphemes[k].morph_token).decode("utf-8")
+
+                        msi = kb_vocab.morpheme_slot_vocab['<UNK>']
+                        if MORPHEME_SLOT in kb_vocab.morpheme_slot_vocab.keys():
+                            msi = kb_vocab.morpheme_slot_vocab[MORPHEME_SLOT]
+
+                        mi = kb_vocab.affix_vocab['<UNK>']
+                        if MORPHEME in kb_vocab.affix_vocab.keys():
+                            mi = kb_vocab.affix_vocab[MORPHEME]
+
+                        ptoken.append_morpheme(msi, mi)
+            parsed_tokens.append(ptoken)
+    lib.release_sentence(sentences, num_sent)
+
+    return parsed_tokens
+
+def process_parsed_sentence(args, parsed_tokens_list: List[ParsedToken], add_cls, kv : KBVocab, affix_set_vocab : AffixSetVocab, rel_pos_dict, rel_pos_dmax):
+    # Input data
+    pos_tags = []
+    stems = []
+    afsets = [] if args.use_afsets else None
+    affixes = []
+    tokens_lengths = []
+
+    # Predicted data
+    predicted_stems = []
+    predicted_afsets = [] if args.use_afsets else None
+    predicted_affixes = [] if args.predict_affixes else None
+
+    predicted_tokens_idx = []
+    predicted_tokens_affixes_idx = [] if args.predict_affixes else None
+    predicted_tokens_affixes_lengths = [] if args.predict_affixes else None
+    # Add <CLS> Token
+    if add_cls:
+        pos_tags.append(kv.pos_tag_vocab['<CLS>'])
+        stems.append(kv.reduced_stem_vocab['<CLS>'])
+        if args.use_afsets:
+            afsets.append(affix_set_vocab.affix_set_to_idx('<CLS>'))
+        tokens_lengths.append(0)
+
+    if (len(parsed_tokens_list) == 0): # New document
+        pos_tags.append(kv.pos_tag_vocab['<SEP>'])
+        stems.append(kv.reduced_stem_vocab['<SEP>'])
+        if args.use_afsets:
+            afsets.append(affix_set_vocab.affix_set_to_idx('<SEP>'))
+        tokens_lengths.append(0)
+
+    else:
+        for ptoken in parsed_tokens_list:
+            for sidx in ptoken.stem_idx:
+                unchanged = True
+                predict = False
+                rval = random.random()
+                if (rval <= 0.15): # 15% of tokens are predicted
+                    predict = True
+                    rval /= 0.15
+                    if(rval < 0.8): # 80% of predicted tokens are masked
+                        unchanged = False
+                        pos_tags.append(kv.pos_tag_vocab['<MSK>'])
+                        stems.append(kv.reduced_stem_vocab['<MSK>'])
+                        if args.use_afsets:
+                            afsets.append(affix_set_vocab.affix_set_to_idx('<MSK>'))
+
+                        vv = rval/0.8
+                        if vv < 0.3: # Include Affixes for 30% of the time to enforce morphology learning
+                            affixes.extend([(v) for v in ptoken.affixes_idx])
+                            tokens_lengths.append(len(ptoken.affixes_idx))
+                        else:
+                            tokens_lengths.append(0)
+                    elif (rval < 0.9): # 10% are replaced by random tokens, 10% are left unchanged
+                        unchanged = False
+                        rnd_pos = random.randint(kv.pos_tag_vocab['<UNK>'], len(kv.pos_tag_vocab)-1)
+                        rnd_stem = random.randint(kv.reduced_stem_vocab['<UNK>'], len(kv.reduced_stem_vocab)-1)
+
+                        pos_tags.append(rnd_pos)
+                        stems.append(rnd_stem)
+                        if args.use_afsets:
+                            afsets.append(affix_set_vocab.random_idx())
+                        elif args.inference_model_file is not None:
+                            affix_set_vocab.random_idx()
+                        vv = rval/0.8
+                        if vv < 0.3: # Include Affixes for 30% of the time to enforce morphology learning
+                            affixes.extend([(v) for v in ptoken.affixes_idx])
+                            tokens_lengths.append(len(ptoken.affixes_idx))
+                        else:
+                            tokens_lengths.append(0)
+                if(unchanged):
+                    pos_tags.append(ptoken.pos_tag_idx)
+                    stems.append(kv.mapped_stem_vocab_idx[sidx])
+                    if args.use_afsets:
+                        afsets.append(affix_set_vocab.affix_set_to_idx(ptoken.affix_set_key()))
+
+                    affixes.extend([(v) for v in ptoken.affixes_idx])
+                    tokens_lengths.append(len(ptoken.affixes_idx))
+
+                if(predict):
+                    predicted_stems.append(kv.mapped_stem_vocab_idx[sidx])
+                    predicted_tokens_idx.append(len(tokens_lengths) - 1)
+                    if args.use_afsets:
+                        predicted_afsets.append(affix_set_vocab.affix_set_to_idx(ptoken.affix_set_key()))
+                    if args.predict_affixes:
+                        predicted_affixes.extend([(v) for v in ptoken.affixes_idx])
+                        if(len(ptoken.affixes_idx) > 0):
+                            predicted_tokens_affixes_idx.append(len(predicted_tokens_idx) - 1)
+                            predicted_tokens_affixes_lengths.append(len(ptoken.affixes_idx))
+
+    rel_pos_arr = np.zeros((len(pos_tags), len(pos_tags))).astype(int) if (rel_pos_dict is not None) else None
+    if rel_pos_arr is not None:
+        for i,pi in enumerate(pos_tags):
+            for j,pj in enumerate(pos_tags):
+                d = i - j
+                if ((d != 0) and (d >= -rel_pos_dmax) and (d <= rel_pos_dmax) and ((not add_cls) or ((i != 0) and (j != 0)))):
+                    rel_pos_arr[i, j] = rel_pos_dict[(pi, pj, d)]
+    return (rel_pos_arr,
+            pos_tags,
+            stems,
+            afsets,
+            affixes,
+            tokens_lengths,
+            predicted_stems,
+            predicted_afsets,
+            predicted_affixes,
+            predicted_tokens_idx,
+            predicted_tokens_affixes_idx,
+            predicted_tokens_affixes_lengths)
+
+def gather_replicated_itemized_data(args, corpus_lines, doc_ends, is_corpus_parsed, max_seq_len, start_line, max_batch_items, kb_vocab : KBVocab, affix_set_vocab : AffixSetVocab, bpe: yttm.BPE, rel_pos_dict, rel_pos_dmax, rank=0, bar=None, num_lines=sys.maxsize,shuffle=False,are_parsed_tokens_split=False):
+    itemized_data = []
+    itemized_parsed_tokens = []
+
+    seq_pos_tags = []
+    seq_stems = []
+    seq_afsets = [] if args.use_afsets else None
+    seq_affixes = []
+    seq_tokens_lengths = []
+    seq_predicted_stems = []
+    seq_predicted_afsets = [] if args.use_afsets else None
+    seq_predicted_affixes = [] if args.predict_affixes else None
+    seq_predicted_tokens_idx = []
+    seq_predicted_tokens_affixes_idx = [] if args.predict_affixes else None
+    seq_predicted_tokens_affixes_lengths = [] if args.predict_affixes else None
+
+    seq_rel_pos_arr = np.zeros((max_seq_len, max_seq_len)).astype(int) if (rel_pos_dict is not None) else None
+
+    seq_parsed_tokens = []
+
+    # if (rank==0):
+    #     print(time_now(), 'Gathering itemized input sequence data from line # {}'.format(start_line + 1), flush=True)
+    lcount = 0
+    while (True):
+        lcount += 1
+        add_cls = (len(seq_parsed_tokens) == 0)
+        if add_cls:
+            ptoken = ParsedToken('<CLS>', decode_prob=1.0, tf_idf=0.001, pos_tag_id=kb_vocab.pos_tag_vocab['<CLS>'], stem_ids=[kb_vocab.reduced_stem_vocab['<CLS>']])
+            ptoken.append_morpheme(kb_vocab.morpheme_slot_vocab['<EOT>'], kb_vocab.affix_vocab['<EOT>'])
+            seq_parsed_tokens.append(ptoken)
+
+        line = corpus_lines[start_line % len(corpus_lines)]
+        start_line = (start_line + 1) % len(corpus_lines)
+        parsed_tokens_line = []
+        if (is_corpus_parsed == True):
+            if are_parsed_tokens_split:
+                parsed_tokens_line = line
+            else:
+                splits = line.split('; ')
+                if (len(splits) > 0):
+                    if ((len(splits[0]) > 4) and ('/' in splits[0]) and (':' in splits[0])):
+                        parsed_tokens_line = [ParsedToken('_', parsed_token=t, line_num=(start_line + 1)) for t in splits]
+        else:
+            parsed_tokens_line = parse_raw_text_lines(line, kb_vocab, bpe)
+
+        if (len(parsed_tokens_line) == 0):
+            ptoken = ParsedToken('<SEP>', decode_prob=1.0, tf_idf=0.001, pos_tag_id=kb_vocab.pos_tag_vocab['<SEP>'], stem_ids=[kb_vocab.reduced_stem_vocab['<SEP>']])
+            ptoken.append_morpheme(kb_vocab.morpheme_slot_vocab['<EOT>'], kb_vocab.affix_vocab['<EOT>'])
+            seq_parsed_tokens.append(ptoken)
+        else:
+            seq_parsed_tokens.extend(parsed_tokens_line)
+
+        (rel_pos_arr,
+         pos_tags,
+         stems,
+         afsets,
+         affixes,
+         tokens_lengths,
+         predicted_stems,
+         predicted_afsets,
+         predicted_affixes,
+         predicted_tokens_idx,
+         predicted_tokens_affixes_idx,
+         predicted_tokens_affixes_lengths) = process_parsed_sentence(args, parsed_tokens_line, add_cls, kb_vocab, affix_set_vocab, rel_pos_dict, rel_pos_dmax)
+        if (len(seq_tokens_lengths) + len(tokens_lengths)) > max_seq_len:
+            data_item = (max_seq_len,
+                         seq_rel_pos_arr,
+                         seq_pos_tags,
+                         seq_stems,
+                         seq_afsets,
+                         seq_affixes,
+                         seq_tokens_lengths,
+                         seq_predicted_stems,
+                         seq_predicted_afsets,
+                         seq_predicted_affixes,
+                         seq_predicted_tokens_idx,
+                         seq_predicted_tokens_affixes_idx,
+                         seq_predicted_tokens_affixes_lengths)
+            itemized_data.append(data_item)
+            itemized_parsed_tokens.append(seq_parsed_tokens)
+            if (shuffle):
+                dcx = random.randint(0, len(doc_ends) - 1) % len(doc_ends)
+                start_line = (doc_ends[dcx]+1) % len(corpus_lines)
+                if random.random() < 0.8: # 80% of the time, start from anywhere within the corpus.
+                    start_line = random.randint(0, len(corpus_lines) - 1) % len(corpus_lines)
+
+            if (len(itemized_data) >= max_batch_items):
+                if ((rank == 0) and (bar is not None)):
+                    bar.update(len(itemized_data))
+                    sys.stdout.flush()
+                return itemized_data, itemized_parsed_tokens, start_line
+
+            if(lcount >= num_lines):
+                return itemized_data, itemized_parsed_tokens, start_line
+
+            seq_pos_tags = []
+            seq_stems = []
+            seq_afsets = [] if args.use_afsets else None
+            seq_affixes = []
+            seq_tokens_lengths = []
+            seq_predicted_stems = []
+            seq_predicted_afsets = [] if args.use_afsets else None
+            seq_predicted_affixes = [] if args.predict_affixes else None
+            seq_predicted_tokens_idx = []
+            seq_predicted_tokens_affixes_idx = [] if args.predict_affixes else None
+            seq_predicted_tokens_affixes_lengths = [] if args.predict_affixes else None
+
+            seq_rel_pos_arr = np.zeros((max_seq_len, max_seq_len)).astype(int) if (rel_pos_dict is not None) else None
+
+            seq_parsed_tokens = []
+
+            if ((rank == 0) and ((len(itemized_data) % (math.floor(0.1 * max_batch_items) + 1)) == 0) and (bar is not None)):
+                bar.update(len(itemized_data))
+                sys.stdout.flush()
+        else:
+            if args.predict_affixes:
+                seq_predicted_tokens_affixes_idx.extend([len(seq_predicted_tokens_idx) + idx for idx in predicted_tokens_affixes_idx])
+
+            seq_predicted_tokens_idx.extend([len(seq_tokens_lengths)+idx for idx in predicted_tokens_idx])
+
+            lt = len(seq_tokens_lengths)
+            if seq_rel_pos_arr is not None:
+                seq_rel_pos_arr[lt:(lt+rel_pos_arr.shape[0]), lt:(lt+rel_pos_arr.shape[1])] = rel_pos_arr
+
+            seq_pos_tags.extend(pos_tags)
+            seq_stems.extend(stems)
+            if args.use_afsets:
+                seq_afsets.extend(afsets)
+            seq_affixes.extend(affixes)
+            seq_tokens_lengths.extend(tokens_lengths)
+            seq_predicted_stems.extend(predicted_stems)
+            if args.use_afsets:
+                seq_predicted_afsets.extend(predicted_afsets)
+            if args.predict_affixes:
+                seq_predicted_affixes.extend(predicted_affixes)
+                seq_predicted_tokens_affixes_lengths.extend(predicted_tokens_affixes_lengths)
+
+            if(lcount >= num_lines):
+                data_item = (max_seq_len,
+                             seq_rel_pos_arr,
+                             seq_pos_tags,
+                             seq_stems,
+                             seq_afsets,
+                             seq_affixes,
+                             seq_tokens_lengths,
+                             seq_predicted_stems,
+                             seq_predicted_afsets,
+                             seq_predicted_affixes,
+                             seq_predicted_tokens_idx,
+                             seq_predicted_tokens_affixes_idx,
+                             seq_predicted_tokens_affixes_lengths)
+                itemized_data.append(data_item)
+                itemized_parsed_tokens.append(seq_parsed_tokens)
+                return itemized_data, itemized_parsed_tokens, start_line
+
+def morpho_seq_collate_wrapper(batch_items):
+    batch_pos_tags = []
+    batch_stems = []
+    batch_afsets = []
+    batch_affixes = []
+    batch_tokens_lengths = []
+    batch_predicted_stems = []
+    batch_predicted_afsets = []
+    batch_predicted_affixes = []
+    batch_predicted_tokens_idx = []
+    batch_predicted_tokens_affixes_idx = []
+    batch_predicted_tokens_affixes_lengths = []
+
+    batch_input_sequence_lengths = []
+
+    max_sequence_len = batch_items[0][0]
+    first_seq_rel_pos_arr = batch_items[0][1]
+
+    batch_rel_pos_arr = np.zeros((len(batch_items), max_sequence_len, max_sequence_len)).astype(int) if (first_seq_rel_pos_arr is not None) else None
+
+    for bidx,data_item in enumerate(batch_items):
+        (max_seq_len,
+         seq_rel_pos_arr,
+         seq_pos_tags,
+         seq_stems,
+         seq_afsets,
+         seq_affixes,
+         seq_tokens_lengths,
+         seq_predicted_stems,
+         seq_predicted_afsets,
+         seq_predicted_affixes,
+         seq_predicted_tokens_idx,
+         seq_predicted_tokens_affixes_idx,
+         seq_predicted_tokens_affixes_lengths) = data_item
+
+        if batch_rel_pos_arr is not None:
+            batch_rel_pos_arr[bidx,:,:] = seq_rel_pos_arr
+        # Need to offset from total number of predicted token indices
+        if seq_predicted_tokens_affixes_idx is not None:
+            batch_predicted_tokens_affixes_idx.extend([(len(batch_predicted_tokens_idx) + t) for t in seq_predicted_tokens_affixes_idx])
+
+        batch_predicted_tokens_idx.extend([(t, len(batch_input_sequence_lengths)) for t in seq_predicted_tokens_idx])
+
+        batch_pos_tags.extend(seq_pos_tags)
+        batch_stems.extend(seq_stems)
+        if seq_afsets is not None:
+            batch_afsets.extend(seq_afsets)
+        batch_affixes.extend(seq_affixes)
+        batch_tokens_lengths.extend(seq_tokens_lengths)
+        batch_predicted_stems.extend(seq_predicted_stems)
+        if seq_predicted_afsets is not None:
+            batch_predicted_afsets.extend(seq_predicted_afsets)
+        if seq_predicted_affixes is not None:
+            batch_predicted_affixes.extend(seq_predicted_affixes)
+            batch_predicted_tokens_affixes_lengths.extend(seq_predicted_tokens_affixes_lengths)
+
+        batch_input_sequence_lengths.append(len(seq_tokens_lengths))
+
+    data_item = (batch_input_sequence_lengths,
+                 batch_rel_pos_arr,
+                 batch_pos_tags,
+                 batch_stems,
+                 batch_afsets,
+                 batch_affixes,
+                 batch_tokens_lengths,
+                 batch_predicted_stems,
+                 batch_predicted_afsets,
+                 batch_predicted_affixes,
+                 batch_predicted_tokens_idx,
+                 batch_predicted_tokens_affixes_idx,
+                 batch_predicted_tokens_affixes_lengths)
+    return data_item
+
+class KBCorpusDataset(Dataset):
+
+    def __init__(self,args,
+                 kb_vocab : KBVocab, affix_set_vocab : AffixSetVocab, bpe_encoder: yttm.BPE,
+                 rel_pos_dict, rel_pos_dmax,
+                 parsed_corpus_lines, doc_ends,  is_corpus_parsed,
+                 start_line, max_batch_items,
+                 max_seq_len = 512,
+                 rank = 0):
+        self.max_seq_len = max_seq_len
+        self.start_line = start_line
+        self.max_batch_items = max_batch_items
+        if (rank==0):
+            with progressbar.ProgressBar(max_value=max_batch_items, redirect_stdout=True) as bar:
+                self.itemized_data, self.itemized_parsed_tokens, self.start_line = gather_replicated_itemized_data(args, parsed_corpus_lines, doc_ends,
+                                                                                                                    is_corpus_parsed,
+                                                                                                                    self.max_seq_len,
+                                                                                                                    self.start_line,
+                                                                                                                    self.max_batch_items,
+                                                                                                                    kb_vocab, affix_set_vocab, bpe_encoder,
+                                                                                                                    rel_pos_dict, rel_pos_dmax,
+                                                                                                                    rank = rank, bar=bar,
+                                                                                                                    shuffle = True)
+        else:
+            self.itemized_data, self.itemized_parsed_tokens, self.start_line = gather_replicated_itemized_data(args, parsed_corpus_lines, doc_ends,
+                                                                                                                is_corpus_parsed,
+                                                                                                                self.max_seq_len,
+                                                                                                                self.start_line,
+                                                                                                                self.max_batch_items,
+                                                                                                                kb_vocab, affix_set_vocab, bpe_encoder,
+                                                                                                                rel_pos_dict, rel_pos_dmax,
+                                                                                                                rank=rank, bar=None,
+                                                                                                                shuffle = True)
+    def __len__(self):
+        return len(self.itemized_data)
+
+    def __getitem__(self, idx):
+        return self.itemized_data[idx]
+
+from itertools import accumulate
+
+from morpho_model import KinyaBERT
+
+def morpho_model_forward(args, data_item, model : KinyaBERT, device, tot_num_affixes):
+    (batch_input_sequence_lengths,
+     batch_rel_pos_arr,
+     batch_pos_tags,
+     batch_stems,
+     batch_afsets,
+     batch_affixes,
+     batch_tokens_lengths,
+     batch_predicted_stems,
+     batch_predicted_afsets,
+     batch_predicted_affixes,
+     batch_predicted_tokens_idx,
+     batch_predicted_tokens_affixes_idx,
+     batch_predicted_tokens_affixes_lengths) = data_item
+
+    tokens_lengths = batch_tokens_lengths  # torch.tensor(batch_tokens_lengths).to(device)
+    input_sequence_lengths = batch_input_sequence_lengths  # torch.tensor(batch_input_sequence_lengths).to(device)
+    pos_tags = torch.tensor(batch_pos_tags).to(device)
+    stems = torch.tensor(batch_stems).to(device)
+    afsets = torch.tensor(batch_afsets).to(device) if args.use_afsets else None
+    affixes = torch.tensor(batch_affixes).to(device)
+
+    predicted_tokens_idx = torch.tensor([s * max(batch_input_sequence_lengths) + t for t, s in batch_predicted_tokens_idx]).to(device)
+    predicted_tokens_affixes_idx = torch.tensor(batch_predicted_tokens_affixes_idx).to(device) if args.predict_affixes else None
+
+    predicted_affixes_prob = None
+    if args.predict_affixes:
+        pred_affixes_list = [batch_predicted_affixes[x - y: x] for x, y in zip(accumulate(batch_predicted_tokens_affixes_lengths), batch_predicted_tokens_affixes_lengths)]
+        afx_prob = torch.zeros(len(pred_affixes_list), tot_num_affixes)
+        for i,lst in enumerate(pred_affixes_list):
+            assert (len(lst) > 0)
+            afx_prob[i,lst] = 1 / len(lst)
+        predicted_affixes_prob = afx_prob.to(device)
+
+    predicted_stems = torch.tensor(batch_predicted_stems).to(device)
+    predicted_afsets = torch.tensor(batch_predicted_afsets).to(device) if args.use_afsets else None
+    rel_pos_arr = torch.from_numpy(batch_rel_pos_arr).to(device) if (batch_rel_pos_arr is not None) else None
+
+    return model(args, rel_pos_arr, tokens_lengths, input_sequence_lengths, pos_tags, stems, afsets, affixes,
+                 predicted_tokens_idx,
+                 predicted_tokens_affixes_idx,
+                 predicted_stems,
+                 predicted_afsets,
+                 predicted_affixes_prob)
+
+def morpho_model_seq_predict(args, data_item, model : KinyaBERT, device,
+                     max_predict_affixes, proposed_stem_ids=None):
+    (max_seq_len,
+     seq_rel_pos_arr,
+     seq_pos_tags,
+     seq_stems,
+     seq_afsets,
+     seq_affixes,
+     seq_tokens_lengths,
+     seq_predicted_stems,
+     seq_predicted_afsets,
+     seq_predicted_affixes,
+     seq_predicted_tokens_idx,
+     seq_predicted_tokens_affixes_idx,
+     seq_predicted_tokens_affixes_lengths) = data_item
+
+    tokens_lengths = seq_tokens_lengths
+    input_sequence_lengths = [len(seq_tokens_lengths)]
+    pos_tags = torch.tensor(seq_pos_tags).to(device)
+    stems = torch.tensor(seq_stems).to(device)
+    afsets = torch.tensor(seq_afsets).to(device) if args.use_afsets else None
+    affixes = torch.tensor(seq_affixes).to(device)
+
+    seq_predicted_token_idx = torch.tensor(seq_predicted_tokens_idx).to(device)
+
+    rel_pos_arr = torch.from_numpy(seq_rel_pos_arr).unsqueeze(0).to(device) if (seq_rel_pos_arr is not None) else None
+
+    return model.predict(args, rel_pos_arr, tokens_lengths, input_sequence_lengths, pos_tags, stems, afsets, affixes,
+                seq_predicted_token_idx,
+                max_predict_affixes, proposed_stem_ids=proposed_stem_ids)

code/morpho_learning_rates.py

@@ -0,0 +1,78 @@
+# coding=utf-8
+# Copyright (c) 2019, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch DataLoader for TFRecords"""
+
+import torch
+
+from torch.optim.lr_scheduler import _LRScheduler
+import math
+
+class AnnealingLR(_LRScheduler):
+    """Anneals the learning rate from start to zero along a cosine curve."""
+
+    DECAY_STYLES = ['linear', 'cosine', 'exponential', 'constant', 'None']
+
+    def __init__(self, optimizer, start_lr, warmup_iter, num_iters, decay_style=None, last_iter=-1):
+        self.optimizer = optimizer
+        self.start_lr = start_lr
+        self.warmup_iter = warmup_iter
+        self.num_iters = last_iter + 1
+        self.end_iter = num_iters
+        self.decay_style = decay_style.lower() if isinstance(decay_style, str) else None
+        self.step(self.num_iters)
+        self.gamma = 0.995
+        self.iters_to_steps = 80.0
+        #print('learning rate decaying', decay_style)
+
+    def get_lr(self):
+        # https://openreview.net/pdf?id=BJYwwY9ll pg. 4
+        if self.warmup_iter > 0 and self.num_iters <= self.warmup_iter:
+            return float(self.start_lr) * self.num_iters / self.warmup_iter
+        else:
+            if self.decay_style == self.DECAY_STYLES[0]:
+                return self.start_lr*((self.end_iter-(self.num_iters-self.warmup_iter))/self.end_iter)
+            elif self.decay_style == self.DECAY_STYLES[1]:
+                return self.start_lr / 2.0 * (math.cos(math.pi * (self.num_iters - self.warmup_iter) / self.end_iter) + 1)
+            elif self.decay_style == self.DECAY_STYLES[2]:
+                #TODO: implement exponential decay
+                return self.start_lr * pow(self.gamma, (self.num_iters - self.warmup_iter) / self.iters_to_steps)
+            else:
+                return self.start_lr
+
+    def step(self, step_num=None):
+        if step_num is None:
+            step_num = self.num_iters + 1
+        self.num_iters = step_num
+        new_lr = self.get_lr()
+        for group in self.optimizer.param_groups:
+            group['lr'] = new_lr
+
+    def state_dict(self):
+        sd = {
+                'start_lr': self.start_lr,
+                'warmup_iter': self.warmup_iter,
+                'num_iters': self.num_iters,
+                'decay_style': self.decay_style,
+                'end_iter': self.end_iter
+        }
+        return sd
+
+    def load_state_dict(self, sd):
+        self.start_lr = sd['start_lr']
+        self.warmup_iter = sd['warmup_iter']
+        self.num_iters = sd['num_iters']
+        self.end_iter = sd['end_iter']
+        self.decay_style = sd['decay_style']
+        self.step(self.num_iters)

code/morpho_model.py

@@ -0,0 +1,850 @@
+# Copyright (c) Antoine Nzeyimana.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import print_function, division
+
+import math
+# Ignore warnings
+import warnings
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.cuda.amp import custom_fwd
+from torch.nn.utils.rnn import pad_sequence
+
+from morpho_transformer import TransformerEncoder, TransformerEncoderLayer, MultiheadAttention
+
+warnings.filterwarnings("ignore")
+
+def gelu(x):
+    """Implementation of the gelu activation function.
+        For information: OpenAI GPT's gelu is slightly different (and gives slightly different results):
+        0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))
+    """
+    return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))
+
+
+def swish(x):
+    return x * torch.sigmoid(x)
+
+def init_bert_params(module):
+    """
+    Initialize the weights specific to the BERT Model.
+    This overrides the default initializations depending on the specified arguments.
+        1. If normal_init_linear_weights is set then weights of linear
+           layer will be initialized using the normal distribution and
+           bais will be set to the specified value.
+        2. If normal_init_embed_weights is set then weights of embedding
+           layer will be initialized using the normal distribution.
+        3. If normal_init_proj_weights is set then weights of
+           in_project_weight for MultiHeadAttention initialized using
+           the normal distribution (to be validated).
+    """
+
+    if isinstance(module, nn.Linear):
+        module.weight.data.normal_(mean=0.0, std=0.02)
+        if module.bias is not None:
+            module.bias.data.zero_()
+    if isinstance(module, nn.Embedding):
+        module.weight.data.normal_(mean=0.0, std=0.02)
+        if module.padding_idx is not None:
+            module.weight.data[module.padding_idx].zero_()
+    if isinstance(module, MultiheadAttention):
+        module.in_proj_weight.data.normal_(mean=0.0, std=0.02)
+
+# From: https://github.com/guolinke/TUPE/blob/master/fairseq/modules/transformer_sentence_encoder.py
+# this is from T5
+def tupe_relative_position_bucket(relative_position, bidirectional=True, num_buckets=32, max_distance=128):
+    ret = 0
+    n = -relative_position
+    if bidirectional:
+        num_buckets //= 2
+        ret += (n < 0).to(torch.long) * num_buckets  # mtf.to_int32(mtf.less(n, 0)) * num_buckets
+        n = torch.abs(n)
+    else:
+        n = torch.max(n, torch.zeros_like(n))
+    # now n is in the range [0, inf)
+
+    # half of the buckets are for exact increments in positions
+    max_exact = num_buckets // 2
+    is_small = n < max_exact
+
+    # The other half of the buckets are for logarithmically bigger bins in positions up to max_distance
+    val_if_large = max_exact + (
+        torch.log(n.float() / max_exact) / math.log(max_distance / max_exact) * (num_buckets - max_exact)
+    ).to(torch.long)
+    val_if_large = torch.min(val_if_large, torch.full_like(val_if_large, num_buckets - 1))
+
+    ret += torch.where(is_small, n, val_if_large)
+    return ret
+
+class BertLayerNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-12):
+        """Construct a layernorm module in the TF style (epsilon inside the square root).
+        """
+        super(BertLayerNorm, self).__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.bias = nn.Parameter(torch.zeros(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, x):
+        u = x.mean(-1, keepdim=True)
+        s = (x - u).pow(2).mean(-1, keepdim=True)
+        x = (x - u) / torch.sqrt(s + self.variance_epsilon)
+        return self.weight * x + self.bias
+
+class BertHeadTransform(nn.Module):
+    def __init__(self, tr_d_model, cls_ctxt_size, layernorm_epsilon):
+        super(BertHeadTransform, self).__init__()
+        self.dense = nn.Linear(tr_d_model, cls_ctxt_size)
+        self.layerNorm = BertLayerNorm(cls_ctxt_size, eps=layernorm_epsilon)
+
+    def forward(self, hidden_states):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = gelu(hidden_states)
+        hidden_states = self.layerNorm(hidden_states)
+        return hidden_states
+
+class TokenClassificationHead(nn.Module):
+    def __init__(self, input_dim, inner_dim, num_classes, pooler_dropout=0.3):
+        super(TokenClassificationHead, self).__init__()
+        self.input_dim = input_dim
+
+        # Experimental
+        #------------------
+        self.dense = nn.Linear(input_dim, inner_dim)
+        self.layerNorm = torch.nn.LayerNorm(inner_dim)
+        self.activation_fn = torch.tanh
+        self.in_dropout = nn.Dropout(p=pooler_dropout)
+        #------------------
+
+        self.out_dropout = nn.Dropout(p=pooler_dropout)
+        self.out_proj = nn.Linear(inner_dim, num_classes)
+
+        self.apply(init_bert_params)
+
+    @custom_fwd
+    def forward(self, features, input_sequence_lengths):
+        # features.shape = S x N x E
+        # Remove [CLS]
+        # len already includes [CLS] in the sequence length count, so number of normal tokens here is (len-1)
+        inputs = [features[1:len, i, :].contiguous().view(-1, self.input_dim) for i,len in enumerate(input_sequence_lengths)]
+        x = torch.cat(inputs, 0) #  B x E
+
+        # Experimental
+        #------------------
+        x = self.in_dropout(x)
+        x = self.dense(x)
+        x = self.layerNorm(x)
+        x = self.activation_fn(x)
+        #------------------
+
+        x = self.out_dropout(x)
+        x = self.out_proj(x)
+        return x
+
+
+class ClassificationHead(nn.Module):
+    def __init__(self, input_dim, inner_dim, num_classes, pooler_dropout=0.0):
+        super(ClassificationHead, self).__init__()
+        self.input_dim = input_dim
+
+        # Experimental
+        #------------------
+        self.dense = nn.Linear(input_dim, inner_dim)
+        self.layerNorm = torch.nn.LayerNorm(inner_dim)
+        self.activation_fn = torch.tanh
+        self.in_dropout = nn.Dropout(p=pooler_dropout)
+        # ------------------
+
+        self.out_dropout = nn.Dropout(p=pooler_dropout)
+        self.out_proj = nn.Linear(inner_dim, num_classes)
+
+        self.apply(init_bert_params)
+
+    def forward(self, features):
+        # features.shape = S x N x E
+        x = features[0, :, :]  # Take [CLS]
+
+        # Experimental
+        # ------------------
+        x = self.in_dropout(x)
+        x = self.dense(x)
+        x = self.layerNorm(x)
+        x = self.activation_fn(x)
+        # ------------------
+
+        x = self.out_dropout(x)
+        x = self.out_proj(x)
+        return x
+
+class MorphoHeadPredictor(nn.Module):
+    def __init__(self, args, stem_embedding_weights, afset_embedding_weights, affix_embedding_weights, tr_d_model, tr_dropout, layernorm_epsilon):
+        super(MorphoHeadPredictor, self).__init__()
+
+        self.stem_transform = BertHeadTransform(tr_d_model, stem_embedding_weights.size(1), layernorm_epsilon)
+        self.stem_decoder = nn.Linear(stem_embedding_weights.size(1), stem_embedding_weights.size(0), bias=False)
+        self.stem_decoder.weight = stem_embedding_weights
+        self.stem_decoder_bias = nn.Parameter(torch.zeros(stem_embedding_weights.size(0)))
+
+        if args.use_afsets:
+            self.afset_transform = BertHeadTransform(tr_d_model, afset_embedding_weights.size(1), layernorm_epsilon)
+            self.afset_decoder = nn.Linear(afset_embedding_weights.size(1), afset_embedding_weights.size(0), bias=False)
+            self.afset_decoder.weight = afset_embedding_weights
+            self.afset_decoder_bias = nn.Parameter(torch.zeros(afset_embedding_weights.size(0)))
+
+        if args.predict_affixes:
+            self.affix_transform = BertHeadTransform(tr_d_model, affix_embedding_weights.size(1), layernorm_epsilon)
+            self.affix_decoder = nn.Linear(affix_embedding_weights.size(1), affix_embedding_weights.size(0), bias=False)
+            self.affix_decoder.weight = affix_embedding_weights
+            self.affix_decoder_bias = nn.Parameter(torch.zeros(affix_embedding_weights.size(0)))
+
+        self.apply(init_bert_params)
+
+    @custom_fwd
+    def forward(self, args, tr_hidden_state,
+                predicted_tokens_idx,
+                predicted_tokens_affixes_idx,
+                predicted_stems,
+                predicted_afsets,
+                predicted_affixes_prob):
+
+        # print('tr_hidden_state.shape',tr_hidden_state.shape)
+        # 1. Crop together predicted tokens
+        # tr_hidden_state.shape = S x N x E
+        token_hidden_state = tr_hidden_state.permute(1,0,2)
+        # N x S x E
+        # print('token_hidden_state.shape',token_hidden_state.shape)
+        token_hidden_state = token_hidden_state.reshape(-1, token_hidden_state.shape[2])
+        # predicted_state.shape: NS x E or B x E
+        # print('token_hidden_state.shape',token_hidden_state.shape)
+        # print('batch_predicted_token_idx',batch_predicted_token_idx)
+        token_hidden_state = torch.index_select(token_hidden_state, 0, index=predicted_tokens_idx)
+        # predicted_state.shape: B x E
+        stem_predicted_state = self.stem_transform(token_hidden_state)
+        # predicted_state.shape: B x F ==> Only containing the 15% tokens to be predicted
+
+        # 3. Propagate to Stem Prediction
+        stem_scores = self.stem_decoder(stem_predicted_state) + self.stem_decoder_bias
+        stem_scores = F.log_softmax(stem_scores, dim=1)
+        stem_loss_avg = F.nll_loss(stem_scores, predicted_stems)
+
+        loss = torch.tensor(0.0).to(stem_loss_avg.device)
+        loss += stem_loss_avg
+
+        # 3*. Propagate to AffixSet Prediction
+        afset_loss_avg = torch.tensor(0.0).to(stem_loss_avg.device)
+        if args.use_afsets and (predicted_afsets is not None):
+            afset_predicted_state = self.afset_transform(token_hidden_state)
+            afset_scores = self.afset_decoder(afset_predicted_state) + self.afset_decoder_bias
+            afset_scores = F.log_softmax(afset_scores, dim=1)
+            afset_loss_avg = F.nll_loss(afset_scores, predicted_afsets)
+            loss += afset_loss_avg
+
+        # 6. Propagate to Affix prediction
+        affix_loss_avg = torch.tensor(0.0).to(stem_loss_avg.device)
+        if args.predict_affixes and (predicted_tokens_affixes_idx is not None) and (predicted_affixes_prob is not None):
+            if predicted_tokens_affixes_idx.nelement() > 0:
+                affix_hidden_state = torch.index_select(token_hidden_state, 0, index=predicted_tokens_affixes_idx)
+                affix_predicted_state = self.affix_transform(affix_hidden_state)
+                affix_scores = self.affix_decoder(affix_predicted_state) + self.affix_decoder_bias
+                affix_scores = F.log_softmax(affix_scores, dim=1)
+                affix_loss_avg = F.kl_div(affix_scores, predicted_affixes_prob, reduction='batchmean')
+
+        loss += affix_loss_avg
+
+        return loss, stem_loss_avg, afset_loss_avg, affix_loss_avg
+
+    def predict(self, args, tr_hidden_state,
+                seq_predicted_token_idx,
+                max_predict_affixes, proposed_stem_ids=None):
+        # print('tr_hidden_state.shape',tr_hidden_state.shape)
+        # 1. Crop together predicted tokens
+        # tr_hidden_state.shape = S x N x E
+        token_hidden_state = tr_hidden_state.permute(1,0,2)
+        # N x S x E
+        # print('token_hidden_state.shape',token_hidden_state.shape)
+        token_hidden_state = token_hidden_state.reshape(-1, token_hidden_state.shape[2])
+        # predicted_state.shape: NS x E = B x E
+        # print('token_hidden_state.shape',token_hidden_state.shape)
+        # print('batch_predicted_token_idx',batch_predicted_token_idx)
+        token_hidden_state = torch.index_select(token_hidden_state, 0, index=seq_predicted_token_idx)
+        # predicted_state.shape: B x E
+        stem_predicted_state = self.stem_transform(token_hidden_state)
+        # predicted_state.shape: B x F ==> Only containing the 15% tokens to be predicted
+
+        # 3. Propagate to Stem Prediction
+        stem_scores = self.stem_decoder(stem_predicted_state) + self.stem_decoder_bias
+        stem_scores = F.softmax(stem_scores, dim=1)
+        # stem_predictions = stem_scores.argmax(dim=1)
+        if proposed_stem_ids is not None:
+            selection = stem_scores[:, proposed_stem_ids]
+            stem_predictions_prob, pred_ids = selection.max(dim=1)
+            pred_list = pred_ids.tolist()
+            stem_predictions = torch.tensor([proposed_stem_ids[ii] for ii in pred_list])
+        else:
+            stem_predictions_prob, stem_predictions = stem_scores.max(dim=1)
+
+        afset_predictions = None
+        afset_predictions_prob = None
+        if args.use_afsets:
+            afset_predicted_state = self.afset_transform(token_hidden_state)
+            afset_scores = self.afset_decoder(afset_predicted_state) + self.afset_decoder_bias
+            afset_scores = F.softmax(afset_scores, dim=1)
+            afset_predictions_prob, afset_predictions = afset_scores.max(dim=1)
+
+        # 5. Predict affix
+        affix_predictions = None
+        if args.predict_affixes:
+            affix_predicted_state = self.affix_transform(token_hidden_state)
+            affix_scores = self.affix_decoder(affix_predicted_state) + self.affix_decoder_bias
+            affix_scores = F.log_softmax(affix_scores, dim=1)
+            _, top_affixes = torch.topk(affix_scores, max_predict_affixes, dim=1)
+            affix_predictions = []
+            for batch in range(top_affixes.shape[0]):
+                affix_predictions.append(top_affixes[batch].tolist())
+
+        return stem_predictions, stem_predictions_prob, afset_predictions, afset_predictions_prob, affix_predictions
+
+class KinyaBERT_MorphoEncoder(nn.Module):
+    def __init__(self, args,
+                 num_stems, num_afsets, num_pos_tags, num_affixes,
+                 num_pos_aware_rel_pos_dict_size,
+                 num_pos_m_embeddings,
+                 num_stem_m_embeddings,
+                 use_affix_bow_m_embedding,
+                 use_pos_aware_rel_pos_bias,
+                 use_tupe_rel_pos_bias,
+                 max_seq_len = 512,
+                 morpho_dim = 80, stem_dim = 160,
+                 morpho_tr_nhead = 4, morpho_tr_nlayers=4,
+                 morpho_tr_dim_feedforward=512, morpho_tr_dropout=0.1, morpho_tr_activation='gelu',
+                 seq_tr_nhead=8, seq_tr_nlayers=8,
+                 seq_tr_dim_feedforward=2048, seq_tr_dropout=0.1, seq_tr_activation='gelu',
+                 tupe_rel_pos_bins: int = 32,
+                 tupe_max_rel_pos: int = 128):
+        super(KinyaBERT_MorphoEncoder, self).__init__()
+
+        self.seq_tr_nhead = seq_tr_nhead
+        self.max_seq_len = max_seq_len
+        self.tot_num_affixes = num_affixes
+        self.num_pos_m_embeddings = num_pos_m_embeddings
+        self.num_stem_m_embeddings = num_stem_m_embeddings
+        self.use_affix_bow_m_embedding = use_affix_bow_m_embedding
+
+        self.seq_tr_d_model = stem_dim
+        self.morpho_dim = morpho_dim
+        self.attn_scale_factor = 2
+
+        self.tot_morpho_idx = 0
+        if self.num_pos_m_embeddings > 0:
+            self.m1_pos_embedding = nn.Embedding(num_pos_tags, self.morpho_dim, padding_idx=0)
+            self.tot_morpho_idx += 1
+
+        if self.num_pos_m_embeddings > 1:
+            self.m2_pos_embedding = nn.Embedding(num_pos_tags, self.morpho_dim, padding_idx=0)
+            self.tot_morpho_idx += 1
+
+        if self.num_pos_m_embeddings > 2:
+            self.m3_pos_embedding = nn.Embedding(num_pos_tags, self.morpho_dim, padding_idx=0)
+            self.tot_morpho_idx += 1
+
+        if self.num_stem_m_embeddings > 0:
+            self.m_stem_embedding = nn.Embedding(num_stems, self.morpho_dim, padding_idx=0)
+            self.tot_morpho_idx += 1
+
+        if args.use_afsets and (num_afsets > 0):
+            self.m_afset_embedding = nn.Embedding(num_afsets, self.morpho_dim, padding_idx=0)
+            self.tot_morpho_idx += 1
+
+        self.seq_tr_d_model += (self.morpho_dim * self.tot_morpho_idx)
+
+        if self.use_affix_bow_m_embedding:
+            self.seq_tr_d_model += self.morpho_dim
+
+        self.s_stem_embedding = nn.Embedding(num_stems, stem_dim, padding_idx=0)
+
+        if args.use_morpho_encoder and (self.seq_tr_d_model > stem_dim):
+            self.m_affix_embedding = nn.Embedding(num_affixes, self.morpho_dim, padding_idx=0)
+            morpho_encoder_layers = TransformerEncoderLayer(self.morpho_dim, morpho_tr_nhead, dim_feedforward=morpho_tr_dim_feedforward, dropout=morpho_tr_dropout, activation=morpho_tr_activation)
+            self.morpho_transformer_encoder = TransformerEncoder(morpho_encoder_layers, morpho_tr_nlayers)
+
+        sequence_encoder_layers = TransformerEncoderLayer(self.seq_tr_d_model, self.seq_tr_nhead, dim_feedforward=seq_tr_dim_feedforward, dropout=seq_tr_dropout, activation=seq_tr_activation)
+        self.seq_transformer_encoder = TransformerEncoder(sequence_encoder_layers, seq_tr_nlayers)
+
+        self.use_pos_aware_rel_pos_bias = use_pos_aware_rel_pos_bias
+        if self.use_pos_aware_rel_pos_bias:
+            self.rel_pos_embedding = nn.Embedding(num_pos_aware_rel_pos_dict_size+1, self.seq_tr_nhead, padding_idx=0)
+
+        # This is from TUPE
+        self.pos = nn.Embedding(self.max_seq_len + 1, self.seq_tr_d_model)
+        self.pos_q_linear = nn.Linear(self.seq_tr_d_model, self.seq_tr_d_model)
+        self.pos_k_linear = nn.Linear(self.seq_tr_d_model, self.seq_tr_d_model)
+        self.pos_scaling = float(self.seq_tr_d_model / self.seq_tr_nhead * self.attn_scale_factor) ** -0.5
+        self.pos_ln = nn.LayerNorm(self.seq_tr_d_model)
+
+        self.use_tupe_rel_pos_bias = use_tupe_rel_pos_bias
+        if self.use_tupe_rel_pos_bias:
+            assert tupe_rel_pos_bins % 2 == 0
+            self.tupe_rel_pos_bins = tupe_rel_pos_bins
+            self.tupe_max_rel_pos = tupe_max_rel_pos
+            self.relative_attention_bias = nn.Embedding(self.tupe_rel_pos_bins + 1, self.seq_tr_nhead)
+            seq_len = self.max_seq_len
+            context_position = torch.arange(seq_len, dtype=torch.long)[:, None]
+            memory_position = torch.arange(seq_len, dtype=torch.long)[None, :]
+            relative_position = memory_position - context_position
+            self.rp_bucket = tupe_relative_position_bucket(
+                relative_position,
+                num_buckets=self.tupe_rel_pos_bins,
+                max_distance=self.tupe_max_rel_pos
+            )
+            # others to [CLS]
+            self.rp_bucket[:, 0] = self.tupe_rel_pos_bins
+            # [CLS] to others, Note: self.tupe_rel_pos_bins // 2 is not used in relative_position_bucket
+            self.rp_bucket[0, :] = self.tupe_rel_pos_bins // 2
+
+        self.apply(init_bert_params)
+
+    def get_tupe_rel_pos_bias(self, seq_len, device):
+        # Assume the input is ordered. If your input token is permuted, you may need to update this accordingly
+        if self.rp_bucket.device != device:
+            self.rp_bucket = self.rp_bucket.to(device)
+        # Adjusted because final x's shape is L x B X E
+        rp_bucket = self.rp_bucket[:seq_len, :seq_len]
+        values = F.embedding(rp_bucket, self.relative_attention_bias.weight)
+        values = values.permute([2, 0, 1])
+        return values.contiguous()
+
+    def get_pos_aware_rel_pos_bias(self, rel_pos_arr, seq_len):
+        rel_pos_idx = rel_pos_arr[:, :seq_len, :seq_len]
+        rel_pos = self.rel_pos_embedding(rel_pos_idx)
+        # N x L x L x h --> N x h x L x L
+        rel_pos = rel_pos.permute([0, 3, 1, 2])
+        rel_pos = rel_pos.contiguous()
+        rel_pos = rel_pos.reshape(-1, seq_len, seq_len)
+        return rel_pos.contiguous()
+
+    def get_position_attn_bias(self, rel_pos_arr, seq_len, batch_size, device):
+        tupe_rel_pos_bias = self.get_tupe_rel_pos_bias(seq_len, device) if self.use_tupe_rel_pos_bias else None
+        pos_aware_rel_pos_bias = self.get_pos_aware_rel_pos_bias(rel_pos_arr, seq_len) if self.use_pos_aware_rel_pos_bias else None
+
+        # This is from TUPE
+        # https://github.com/guolinke/TUPE/blob/master/fairseq/modules/transformer_sentence_encoder.py
+        # 0 is for other-to-cls 1 is for cls-to-other
+        # Assume the input is ordered. If your input token is permuted, you may need to update this accordingly
+        weight = self.pos_ln(self.pos.weight[:seq_len + 1, :])
+        pos_q =  self.pos_q_linear(weight).view(seq_len + 1, self.seq_tr_nhead, -1).transpose(0, 1) * self.pos_scaling
+        pos_k =  self.pos_k_linear(weight).view(seq_len + 1, self.seq_tr_nhead, -1).transpose(0, 1)
+        abs_pos_bias = torch.bmm(pos_q, pos_k.transpose(1, 2))
+        # p_0 \dot p_0 is cls to others
+        cls_2_other = abs_pos_bias[:, 0, 0]
+        # p_1 \dot p_1 is others to cls
+        other_2_cls = abs_pos_bias[:, 1, 1]
+        # offset
+        abs_pos_bias = abs_pos_bias[:, 1:, 1:]
+        abs_pos_bias[:, :, 0] = other_2_cls.view(-1, 1)
+        abs_pos_bias[:, 0, :] = cls_2_other.view(-1, 1)
+
+        if tupe_rel_pos_bias is not None:
+            abs_pos_bias += tupe_rel_pos_bias
+
+        abs_pos_bias = abs_pos_bias.unsqueeze(0).expand(batch_size, -1, -1, -1).reshape(-1, seq_len, seq_len)
+
+        if pos_aware_rel_pos_bias is not None:
+            abs_pos_bias += pos_aware_rel_pos_bias
+
+        return abs_pos_bias
+
+
+    @custom_fwd
+    def forward(self, args, rel_pos_arr, tokens_lengths, input_sequence_lengths, pos_tags, stems, afsets, affixes):
+        # L = len(tokens_lengths)
+        # affixes: (N)
+        # stems: (L)
+        # pos_tags: (L)
+
+        device = stems.device
+        x_embed = None
+        if self.num_pos_m_embeddings > 0:
+            xm_pos1 = self.m1_pos_embedding(pos_tags)
+            x_embed = torch.unsqueeze(xm_pos1,0)
+
+        if self.num_pos_m_embeddings > 1:
+            xm_pos2 = self.m2_pos_embedding(pos_tags)
+            xm_pos2 = torch.unsqueeze(xm_pos2,0)
+            x_embed = torch.cat((x_embed, xm_pos2), 0)
+
+        if self.num_pos_m_embeddings > 2:
+            xm_pos3 = self.m3_pos_embedding(pos_tags)
+            xm_pos3 = torch.unsqueeze(xm_pos3,0)
+            x_embed = torch.cat((x_embed, xm_pos3), 0)
+
+        if self.num_stem_m_embeddings > 0:
+            xm_stem = self.m_stem_embedding(stems)
+            xm_stem = torch.unsqueeze(xm_stem,0)
+            if x_embed is not None:
+                x_embed = torch.cat((x_embed, xm_stem), 0)
+            else:
+                x_embed = xm_stem
+
+        if (args.use_afsets) and (afsets is not None):
+            xm_afset = self.m_afset_embedding(afsets)
+            xm_afset = torch.unsqueeze(xm_afset,0)
+            if x_embed is not None:
+                x_embed = torch.cat((x_embed, xm_afset), 0)
+            else:
+                x_embed = xm_afset
+        # All above: (1,L,E)
+
+        #x_embed = torch.cat((xm_pos1, xm_pos2, xm_pos3, xm_stem, xm_afset), 0)
+        # x_embed: (4,L,E)
+
+        m_masks_padded = None
+        has_morphemes = False
+        if args.use_morpho_encoder:
+            afx = affixes.split(tokens_lengths)
+            # [[2,4,5], [6,7]]
+            afx_padded = pad_sequence(afx, batch_first=False)
+            # afx_padded: (M,L), M: max morphological length
+            if afx_padded.nelement() > 0:
+                has_morphemes = True
+                xm_affix = self.m_affix_embedding(afx_padded)
+                # xa_embed: (M,L,E)
+                if x_embed is not None:
+                    x_embed = torch.cat((x_embed, xm_affix), 0)
+                else:
+                    x_embed = xm_affix
+
+                m_masks = [torch.zeros((x+(self.tot_morpho_idx)), dtype=torch.bool, device = device) for x in tokens_lengths]
+                m_masks_padded = pad_sequence(m_masks, batch_first=True, padding_value=1) # Shape: (L, M+4)
+
+        # x_embed: (4+M,L,E)
+        morpho_input = None
+        if (x_embed is not None) and args.use_morpho_encoder:
+            morpho_transformer_output = self.morpho_transformer_encoder(x_embed, src_key_padding_mask=m_masks_padded)  # Shape: (M+4, L, E)
+            if (self.tot_morpho_idx > 0) and (self.use_affix_bow_m_embedding):
+                heads = morpho_transformer_output[:(self.tot_morpho_idx), :, :]
+                affixes_bow = torch.sum(morpho_transformer_output[(self.tot_morpho_idx):, :, :], 0, keepdim=True) if has_morphemes else torch.zeros((1, stems.size(0), self.morpho_dim), device = device)
+                morpho_input = torch.cat((heads, affixes_bow), 0)
+            elif (self.tot_morpho_idx > 0):
+                morpho_input = morpho_transformer_output[:(self.tot_morpho_idx), :, :]
+            elif (self.use_affix_bow_m_embedding):
+                morpho_input = torch.sum(morpho_transformer_output[(self.tot_morpho_idx):, :, :], 0, keepdim=True) if has_morphemes else torch.zeros((1, stems.size(0), self.morpho_dim), device = device)
+        elif (self.use_affix_bow_m_embedding):
+            # No POS, no STEM, No afset, No affixes, but allows affixes_bow
+            morpho_input = torch.zeros((1, stems.size(0), self.morpho_dim), device = device)
+
+        input_sequences = self.s_stem_embedding(stems) # (L, E')
+        if morpho_input is not None:
+            # 4 x L x E ==> L x 4 x E
+            #i.e. K x L x E ==> L x K x E, K: number of morpho components from the morpho encoder tier
+            morpho_input = morpho_input.permute(1, 0, 2)
+            L = morpho_input.size(0)
+            morpho_input = morpho_input.contiguous().view(L,-1) # (L, 4E), i.e. (L, KE)
+            input_sequences = torch.cat((morpho_input, input_sequences), 1)
+
+        lists = input_sequences.split(input_sequence_lengths, 0) # len(input_sequence_lengths) = N (i.e. Batch Size, e.g. 32)
+        tr_padded = pad_sequence(lists, batch_first=False)
+
+        seq_len = tr_padded.size(0)
+        batch_size = tr_padded.size(1)
+
+        abs_pos_bias = self.get_position_attn_bias(rel_pos_arr, seq_len, batch_size, device)
+
+        masks = [torch.zeros(x, dtype=torch.bool, device = device) for x in input_sequence_lengths]
+        masks_padded = pad_sequence(masks, batch_first=True, padding_value=1) # Shape: N x S
+
+        transformer_output = self.seq_transformer_encoder(tr_padded, attn_bias = abs_pos_bias, src_key_padding_mask = masks_padded) # Shape: L x N x E, with L = max sequence length
+
+        return transformer_output
+
+
+class KinyaBERT(nn.Module):
+    def __init__(self, args,
+                 num_stems, num_afsets, num_pos_tags, num_affixes,
+                 num_rel_pos_dict_size,
+                 num_pos_m_embeddings,
+                 num_stem_m_embeddings,
+                 use_affix_bow_m_embedding,
+                 use_pos_aware_rel_pos_bias,
+                 use_tupe_rel_pos_bias,
+                 max_seq_len = 512,
+                 morpho_dim=80, stem_dim=160,
+                 morpho_tr_nhead=4, morpho_tr_nlayers=4,
+                 morpho_tr_dim_feedforward=512, morpho_tr_dropout=0.1, morpho_tr_activation='gelu',
+                 seq_tr_nhead=8, seq_tr_nlayers=8,
+                 seq_tr_dim_feedforward=2048, seq_tr_dropout=0.1, seq_tr_activation='gelu',
+                 layernorm_epsilon = 1e-6,
+                 tupe_rel_pos_bins: int = 32,
+                 tupe_max_rel_pos: int = 128):
+        super(KinyaBERT, self).__init__()
+
+        self.encoder = KinyaBERT_MorphoEncoder(args, num_stems, num_afsets, num_pos_tags, num_affixes,
+                 num_rel_pos_dict_size,
+                 num_pos_m_embeddings,
+                 num_stem_m_embeddings,
+                 use_affix_bow_m_embedding,
+                 use_pos_aware_rel_pos_bias,
+                 use_tupe_rel_pos_bias,
+                 max_seq_len = max_seq_len,
+                 morpho_dim = morpho_dim, stem_dim = stem_dim,
+                 morpho_tr_nhead = morpho_tr_nhead, morpho_tr_nlayers=morpho_tr_nlayers,
+                 morpho_tr_dim_feedforward=morpho_tr_dim_feedforward, morpho_tr_dropout=morpho_tr_dropout, morpho_tr_activation=morpho_tr_activation,
+                 seq_tr_nhead=seq_tr_nhead, seq_tr_nlayers=seq_tr_nlayers,
+                 seq_tr_dim_feedforward=seq_tr_dim_feedforward, seq_tr_dropout=seq_tr_dropout, seq_tr_activation=seq_tr_activation,
+                 tupe_rel_pos_bins = tupe_rel_pos_bins,
+                 tupe_max_rel_pos = tupe_max_rel_pos)
+
+        self.predictor = MorphoHeadPredictor(args, self.encoder.s_stem_embedding.weight,
+                                             self.encoder.m_afset_embedding.weight if (num_afsets > 0) else None,
+                                             self.encoder.m_affix_embedding.weight if args.predict_affixes else None,
+                                             self.encoder.seq_tr_d_model, seq_tr_dropout, layernorm_epsilon)
+
+    @custom_fwd
+    def forward(self, args, rel_pos_arr, tokens_lengths, input_sequence_lengths, pos_tags, stems, afsets, affixes,
+                predicted_tokens_idx,
+                predicted_tokens_affixes_idx,
+                predicted_stems,
+                predicted_afsets,
+                predicted_affixes_prob):
+
+        tr_hidden_state = self.encoder(args, rel_pos_arr, tokens_lengths, input_sequence_lengths, pos_tags, stems, afsets, affixes)
+
+        return self.predictor(args, tr_hidden_state,
+                              predicted_tokens_idx,
+                              predicted_tokens_affixes_idx,
+                              predicted_stems,
+                              predicted_afsets,
+                              predicted_affixes_prob)
+
+    def predict(self, args, rel_pos_arr, tokens_lengths, input_sequence_lengths, pos_tags, stems, afsets, affixes,
+                seq_predicted_token_idx,
+                max_predict_affixes, proposed_stem_ids=None):
+
+        tr_hidden_state = self.encoder(args, rel_pos_arr, tokens_lengths, input_sequence_lengths, pos_tags, stems, afsets, affixes)
+
+        return self.predictor.predict(args, tr_hidden_state,
+                seq_predicted_token_idx,
+                max_predict_affixes, proposed_stem_ids=proposed_stem_ids)
+
+class KinyaBERTClassifier(nn.Module):
+    def __init__(self, args,
+                 num_stems, num_afsets, num_pos_tags, num_affixes,
+                 num_rel_pos_dict_size,
+                 num_classes,
+                 num_pos_m_embeddings,
+                 num_stem_m_embeddings,
+                 use_affix_bow_m_embedding,
+                 use_pos_aware_rel_pos_bias,
+                 use_tupe_rel_pos_bias,
+                 max_seq_len = 512,
+                 morpho_dim=80, stem_dim=160,
+                 morpho_tr_nhead=4, morpho_tr_nlayers=4,
+                 morpho_tr_dim_feedforward=512, morpho_tr_dropout=0.1, morpho_tr_activation='gelu',
+                 seq_tr_nhead=8, seq_tr_nlayers=8,
+                 seq_tr_dim_feedforward=2048, seq_tr_dropout=0.1, pooler_dropout=0.0, seq_tr_activation='gelu',
+                 tupe_rel_pos_bins: int = 32,
+                 tupe_max_rel_pos: int = 128):
+        super(KinyaBERTClassifier, self).__init__()
+
+        self.encoder = KinyaBERT_MorphoEncoder(args, num_stems, num_afsets, num_pos_tags, num_affixes,
+                 num_rel_pos_dict_size,
+                 num_pos_m_embeddings,
+                 num_stem_m_embeddings,
+                 use_affix_bow_m_embedding,
+                 use_pos_aware_rel_pos_bias,
+                 use_tupe_rel_pos_bias,
+                 max_seq_len = max_seq_len,
+                 morpho_dim = morpho_dim, stem_dim = stem_dim,
+                 morpho_tr_nhead = morpho_tr_nhead, morpho_tr_nlayers=morpho_tr_nlayers,
+                 morpho_tr_dim_feedforward=morpho_tr_dim_feedforward, morpho_tr_dropout=morpho_tr_dropout, morpho_tr_activation=morpho_tr_activation,
+                 seq_tr_nhead=seq_tr_nhead, seq_tr_nlayers=seq_tr_nlayers,
+                 seq_tr_dim_feedforward=seq_tr_dim_feedforward, seq_tr_dropout=seq_tr_dropout, seq_tr_activation=seq_tr_activation,
+                 tupe_rel_pos_bins = tupe_rel_pos_bins,
+                 tupe_max_rel_pos = tupe_max_rel_pos)
+
+        self.cls_head = ClassificationHead(self.encoder.seq_tr_d_model, num_classes * 32, num_classes, pooler_dropout=pooler_dropout)
+
+    @custom_fwd
+    def forward(self, args, rel_pos_arr, tokens_lengths, input_sequence_lengths, pos_tags, stems, afsets, affixes):
+        tr_hidden_state = self.encoder(args, rel_pos_arr, tokens_lengths, input_sequence_lengths, pos_tags, stems, afsets, affixes)
+        return self.cls_head(tr_hidden_state)
+
+class KinyaBERTSequenceTagger(nn.Module):
+    def __init__(self, args,
+                 num_stems, num_afsets, num_pos_tags, num_affixes,
+                 num_rel_pos_dict_size,
+                 num_classes,
+                 num_pos_m_embeddings,
+                 num_stem_m_embeddings,
+                 use_affix_bow_m_embedding,
+                 use_pos_aware_rel_pos_bias,
+                 use_tupe_rel_pos_bias,
+                 max_seq_len = 512,
+                 morpho_dim=80, stem_dim=160,
+                 morpho_tr_nhead=4, morpho_tr_nlayers=4,
+                 morpho_tr_dim_feedforward=512, morpho_tr_dropout=0.1, morpho_tr_activation='gelu',
+                 seq_tr_nhead=8, seq_tr_nlayers=8,
+                 seq_tr_dim_feedforward=2048, seq_tr_dropout=0.1, pooler_dropout=0.0, seq_tr_activation='gelu',
+                 tupe_rel_pos_bins: int = 32,
+                 tupe_max_rel_pos: int = 128):
+        super(KinyaBERTSequenceTagger, self).__init__()
+
+        self.encoder = KinyaBERT_MorphoEncoder(args, num_stems, num_afsets, num_pos_tags, num_affixes,
+                 num_rel_pos_dict_size,
+                 num_pos_m_embeddings,
+                 num_stem_m_embeddings,
+                 use_affix_bow_m_embedding,
+                 use_pos_aware_rel_pos_bias,
+                 use_tupe_rel_pos_bias,
+                 max_seq_len = max_seq_len,
+                 morpho_dim = morpho_dim, stem_dim = stem_dim,
+                 morpho_tr_nhead = morpho_tr_nhead, morpho_tr_nlayers=morpho_tr_nlayers,
+                 morpho_tr_dim_feedforward=morpho_tr_dim_feedforward, morpho_tr_dropout=morpho_tr_dropout, morpho_tr_activation=morpho_tr_activation,
+                 seq_tr_nhead=seq_tr_nhead, seq_tr_nlayers=seq_tr_nlayers,
+                 seq_tr_dim_feedforward=seq_tr_dim_feedforward, seq_tr_dropout=seq_tr_dropout, seq_tr_activation=seq_tr_activation,
+                 tupe_rel_pos_bins = tupe_rel_pos_bins,
+                 tupe_max_rel_pos = tupe_max_rel_pos)
+
+        self.cls_head = TokenClassificationHead(self.encoder.seq_tr_d_model, num_classes * 32, num_classes, pooler_dropout=pooler_dropout)
+
+    @custom_fwd
+    def forward(self, args, rel_pos_arr, tokens_lengths, input_sequence_lengths, pos_tags, stems, afsets, affixes):
+        tr_hidden_state = self.encoder(args, rel_pos_arr, tokens_lengths, input_sequence_lengths, pos_tags, stems, afsets, affixes)
+        return self.cls_head(tr_hidden_state, input_sequence_lengths)
+
+from morpho_data_loaders import KBVocab, AffixSetVocab
+
+def kinyabert_base(kb_vocab : KBVocab, affix_set_vocab : AffixSetVocab, rel_pos_dict, device, args, saved_model_file = None) -> KinyaBERT:
+    num_pos_tags = len(kb_vocab.pos_tag_vocab) + 1
+    num_stems = len(kb_vocab.reduced_stem_vocab) + 1
+    num_afsets = (len(affix_set_vocab.affix_set_vocab_idx) + 1) if args.use_afsets else 0
+    num_affixes = len(kb_vocab.affix_vocab) + 1
+
+    num_rel_pos_dict_size = (len(rel_pos_dict)+1) if (rel_pos_dict is not None) else 0
+
+    activation_fn = 'gelu'
+
+    kb_model = KinyaBERT(args, num_stems, num_afsets, num_pos_tags, num_affixes,
+                 num_rel_pos_dict_size,
+                 args.num_pos_m_embeddings,
+                 args.num_stem_m_embeddings,
+                 args.use_affix_bow_m_embedding,
+                 args.use_pos_aware_rel_pos_bias,
+                 args.use_tupe_rel_pos_bias,
+                 max_seq_len = args.max_seq_len,
+                 morpho_dim = args.morpho_dim, stem_dim = args.stem_dim,
+                 morpho_tr_nhead = args.morpho_tr_nhead, morpho_tr_nlayers=args.morpho_tr_nlayers,
+                 morpho_tr_dim_feedforward=args.morpho_tr_dim_feedforward, morpho_tr_dropout=args.morpho_tr_dropout, morpho_tr_activation=activation_fn,
+                 seq_tr_nhead=args.seq_tr_nhead, seq_tr_nlayers=args.seq_tr_nlayers,
+                 seq_tr_dim_feedforward=args.seq_tr_dim_feedforward, seq_tr_dropout=args.seq_tr_dropout, seq_tr_activation=activation_fn,
+                 layernorm_epsilon = args.layernorm_epsilon).to(device)
+
+    if saved_model_file is not None:
+        kb_state_dict = torch.load(saved_model_file, map_location=device)
+        kb_model.load_state_dict(kb_state_dict['model_state_dict'])
+
+    return kb_model
+
+def kinyabert_base_classifier(num_classes, kb_vocab : KBVocab, affix_set_vocab : AffixSetVocab, rel_pos_dict, device, args, saved_model_file = None, pooler_dropout=0.0) -> KinyaBERTClassifier:
+    num_pos_tags = len(kb_vocab.pos_tag_vocab) + 1
+    num_stems = len(kb_vocab.reduced_stem_vocab) + 1
+    num_afsets = (len(affix_set_vocab.affix_set_vocab_idx) + 1) if args.use_afsets else 0
+    num_affixes = len(kb_vocab.affix_vocab) + 1
+
+    num_rel_pos_dict_size = (len(rel_pos_dict)+1) if (rel_pos_dict is not None) else 0
+
+    activation_fn = 'gelu'
+
+    kb_model = KinyaBERTClassifier(args, num_stems, num_afsets, num_pos_tags, num_affixes,
+                 num_rel_pos_dict_size,
+                 num_classes,
+                 args.num_pos_m_embeddings,
+                 args.num_stem_m_embeddings,
+                 args.use_affix_bow_m_embedding,
+                 args.use_pos_aware_rel_pos_bias,
+                 args.use_tupe_rel_pos_bias,
+                 max_seq_len = args.max_seq_len,
+                 morpho_dim = args.morpho_dim, stem_dim = args.stem_dim,
+                 morpho_tr_nhead = args.morpho_tr_nhead, morpho_tr_nlayers=args.morpho_tr_nlayers,
+                 morpho_tr_dim_feedforward=args.morpho_tr_dim_feedforward, morpho_tr_dropout=args.morpho_tr_dropout, morpho_tr_activation=activation_fn,
+                 seq_tr_nhead=args.seq_tr_nhead, seq_tr_nlayers=args.seq_tr_nlayers,
+                 seq_tr_dim_feedforward=args.seq_tr_dim_feedforward, seq_tr_dropout=args.seq_tr_dropout, seq_tr_activation=activation_fn,
+                 pooler_dropout = pooler_dropout).to(device)
+
+    if saved_model_file is not None:
+        kb_state_dict = torch.load(saved_model_file, map_location=device)
+        kb_model.load_state_dict(kb_state_dict['model_state_dict'])
+
+    return kb_model
+
+def kinyabert_base_classifier_from_pretrained(num_classes, kb_vocab : KBVocab, affix_set_vocab : AffixSetVocab, rel_pos_dict, device, args, pretrained_model_file, ddp = False, pooler_dropout=0.0)\
+        -> KinyaBERTClassifier:
+    from torch.nn.parallel import DistributedDataParallel as DDP
+    pretrained_model = kinyabert_base(kb_vocab, affix_set_vocab, rel_pos_dict, device, args)
+    classifier_model = kinyabert_base_classifier(num_classes, kb_vocab, affix_set_vocab, rel_pos_dict, device, args, pooler_dropout=pooler_dropout)
+
+    kb_state_dict = torch.load(pretrained_model_file, map_location=device)
+    if ddp:
+        ddp_model = DDP(pretrained_model)
+        ddp_model.load_state_dict(kb_state_dict['model_state_dict'])
+        pretrained_model = ddp_model.module
+    else:
+        pretrained_model.load_state_dict(kb_state_dict['model_state_dict'])
+
+    classifier_model.encoder.load_state_dict(pretrained_model.encoder.state_dict())
+
+    return classifier_model
+
+def kinyabert_base_sequence_tagger(num_classes, kb_vocab : KBVocab, affix_set_vocab : AffixSetVocab, rel_pos_dict, device, args, saved_model_file = None, pooler_dropout=0.0)\
+        -> KinyaBERTSequenceTagger:
+    num_pos_tags = len(kb_vocab.pos_tag_vocab) + 1
+    num_stems = len(kb_vocab.reduced_stem_vocab) + 1
+    num_afsets = (len(affix_set_vocab.affix_set_vocab_idx) + 1) if args.use_afsets else 0
+    num_affixes = len(kb_vocab.affix_vocab) + 1
+
+    num_rel_pos_dict_size = (len(rel_pos_dict)+1) if (rel_pos_dict is not None) else 0
+
+    activation_fn = 'gelu'
+
+    kb_model = KinyaBERTSequenceTagger(args, num_stems, num_afsets, num_pos_tags, num_affixes,
+                 num_rel_pos_dict_size,
+                 num_classes,
+                 args.num_pos_m_embeddings,
+                 args.num_stem_m_embeddings,
+                 args.use_affix_bow_m_embedding,
+                 args.use_pos_aware_rel_pos_bias,
+                 args.use_tupe_rel_pos_bias,
+                 max_seq_len = args.max_seq_len,
+                 morpho_dim = args.morpho_dim, stem_dim = args.stem_dim,
+                 morpho_tr_nhead = args.morpho_tr_nhead, morpho_tr_nlayers=args.morpho_tr_nlayers,
+                 morpho_tr_dim_feedforward=args.morpho_tr_dim_feedforward, morpho_tr_dropout=args.morpho_tr_dropout, morpho_tr_activation=activation_fn,
+                 seq_tr_nhead=args.seq_tr_nhead, seq_tr_nlayers=args.seq_tr_nlayers,
+                 seq_tr_dim_feedforward=args.seq_tr_dim_feedforward, seq_tr_dropout=args.seq_tr_dropout, pooler_dropout=pooler_dropout, seq_tr_activation=activation_fn).to(device)
+
+    if saved_model_file is not None:
+        kb_state_dict = torch.load(saved_model_file, map_location=device)
+        kb_model.load_state_dict(kb_state_dict['model_state_dict'])
+
+    return kb_model
+
+def kinyabert_base_tagger_from_pretrained(num_classes, kb_vocab : KBVocab, affix_set_vocab : AffixSetVocab, rel_pos_dict, device, args, pretrained_model_file, ddp = False, pooler_dropout=0.0)\
+        -> KinyaBERTSequenceTagger:
+    from torch.nn.parallel import DistributedDataParallel as DDP
+    pretrained_model = kinyabert_base(kb_vocab, affix_set_vocab, rel_pos_dict, device, args)
+    sequence_tagger_model = kinyabert_base_sequence_tagger(num_classes, kb_vocab, affix_set_vocab, rel_pos_dict, device, args, pooler_dropout=pooler_dropout)
+
+    kb_state_dict = torch.load(pretrained_model_file, map_location=device)
+    if ddp:
+        ddp_model = DDP(pretrained_model)
+        ddp_model.load_state_dict(kb_state_dict['model_state_dict'])
+        pretrained_model = ddp_model.module
+    else:
+        pretrained_model.load_state_dict(kb_state_dict['model_state_dict'])
+
+    sequence_tagger_model.encoder.load_state_dict(pretrained_model.encoder.state_dict())
+
+    return sequence_tagger_model

code/morpho_stub.py

@@ -0,0 +1,82 @@
+# TabuLM — morphological analyzer stub
+# Drop-in replacement for parse_raw_text_lines that works without libkinlp.so.
+# Uses regex heuristics to assign word types and falls back to BPE for all stems.
+
+import re
+from typing import List
+
+import youtokentome as yttm
+
+_NUMERIC_RE = re.compile(
+    r'^[\d,.\s]+(%|Frw|RWF|km|kg|ha|m|m²|L|MW|USD|acres|ha)?$', re.IGNORECASE
+)
+_DATE_RE = re.compile(
+    r'^\d{4}(-\d{2}(-\d{2})?)?$|^\d{1,2}/\d{1,2}/\d{2,4}$'
+)
+_SPECIAL_TOKEN_RE = re.compile(r'^\[.*\]$')
+_UPPER_START_RE = re.compile(r'^[A-Z][a-z]')
+
+
+def _classify_word_type(word: str) -> str:
+    """Return the KBVocab word-type prefix for a given surface form."""
+    if _SPECIAL_TOKEN_RE.match(word):
+        return 'T'
+    if _DATE_RE.match(word) or _NUMERIC_RE.match(word):
+        return 'NU'
+    if _UPPER_START_RE.match(word):
+        return 'NP'
+    return 'T'
+
+
+def parse_text_stub(text: str, kb_vocab, bpe: yttm.BPE) -> List:
+    """
+    Stub replacement for parse_raw_text_lines from morpho_data_loaders.
+    Takes a space-separated text string and returns a list of ParsedToken objects.
+    No libkinlp.so / CFFI required — all stems resolved via BPE fallback.
+    """
+    from morpho_data_loaders import ParsedToken
+
+    unk_pos = kb_vocab.pos_tag_vocab.get('<UNK>', 1)
+    unk_stem = kb_vocab._stem_vocab.get('<UNK>', 1)
+
+    parsed_tokens: List[ParsedToken] = []
+
+    for word in text.split():
+        if not word:
+            continue
+
+        word_type = _classify_word_type(word)
+
+        if word_type == 'NU':
+            stem_key = f'NU:{word}'
+            si = kb_vocab._stem_vocab.get(stem_key, unk_stem)
+            ptoken = ParsedToken(
+                word, decode_prob=1.0, tf_idf=0.001,
+                pos_tag_id=unk_pos, stem_ids=[si]
+            )
+            parsed_tokens.append(ptoken)
+            continue
+
+        try:
+            subwords = bpe.encode(word, output_type=yttm.OutputType.SUBWORD)
+        except Exception:
+            subwords = []
+
+        if not subwords:
+            stem_key = f'{word_type}:{word}'
+            si = kb_vocab._stem_vocab.get(stem_key, unk_stem)
+            subwords_sids = [si]
+        else:
+            subwords_sids = []
+            for sw in subwords:
+                stem_key = f'{word_type}:{sw}'
+                si = kb_vocab._stem_vocab.get(stem_key, unk_stem)
+                subwords_sids.append(si)
+
+        ptoken = ParsedToken(
+            word, decode_prob=1.0, tf_idf=0.001,
+            pos_tag_id=unk_pos, stem_ids=subwords_sids
+        )
+        parsed_tokens.append(ptoken)
+
+    return parsed_tokens

code/morpho_transformer.py

@@ -0,0 +1,642 @@
+# Copyright (c) Antoine Nzeyimana.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+# MIT License
+#
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# 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.
+
+r"""Functional interface"""
+from typing import Callable, List, Optional, Tuple
+import warnings
+
+import torch
+
+from torch import Tensor
+
+from torch.nn.init import xavier_uniform_
+from torch.nn.init import constant_
+from torch.nn.init import xavier_normal_
+from torch.nn.parameter import Parameter
+from torch.nn import Module, Linear, LayerNorm, Dropout, ModuleList
+import copy
+from typing import Optional, Any
+
+import torch.nn.functional as F
+from torch.nn.functional import linear, pad, softmax, dropout
+
+def multi_head_attention_forward(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    embed_dim_to_check: int,
+    num_heads: int,
+    in_proj_weight: Tensor,
+    in_proj_bias: Optional[Tensor],
+    bias_k: Optional[Tensor],
+    bias_v: Optional[Tensor],
+    add_zero_attn: bool,
+    dropout_p: float,
+    out_proj_weight: Tensor,
+    out_proj_bias: Optional[Tensor],
+    training: bool = True,
+    key_padding_mask: Optional[Tensor] = None,
+    need_weights: bool = True,
+    attn_mask: Optional[Tensor] = None,
+    use_separate_proj_weight: bool = False,
+    q_proj_weight: Optional[Tensor] = None,
+    k_proj_weight: Optional[Tensor] = None,
+    v_proj_weight: Optional[Tensor] = None,
+    static_k: Optional[Tensor] = None,
+    static_v: Optional[Tensor] = None,
+    attn_bias: Optional[Tensor] = None,
+    scale_factor = 1,
+) -> Tuple[Tensor, Optional[Tensor]]:
+    r"""
+    Args:
+        query, key, value: map a query and a set of key-value pairs to an output.
+            See "Attention Is All You Need" for more details.
+        embed_dim_to_check: total dimension of the model.
+        num_heads: parallel attention heads.
+        in_proj_weight, in_proj_bias: input projection weight and bias.
+        bias_k, bias_v: bias of the key and value sequences to be added at dim=0.
+        add_zero_attn: add a new batch of zeros to the key and
+                       value sequences at dim=1.
+        dropout_p: probability of an element to be zeroed.
+        out_proj_weight, out_proj_bias: the output projection weight and bias.
+        training: apply dropout if is ``True``.
+        key_padding_mask: if provided, specified padding elements in the key will
+            be ignored by the attention. This is an binary mask. When the value is True,
+            the corresponding value on the attention layer will be filled with -inf.
+        need_weights: output attn_output_weights.
+        attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
+            the batches while a 3D mask allows to specify a different mask for the entries of each batch.
+        use_separate_proj_weight: the function accept the proj. weights for query, key,
+            and value in different forms. If false, in_proj_weight will be used, which is
+            a combination of q_proj_weight, k_proj_weight, v_proj_weight.
+        q_proj_weight, k_proj_weight, v_proj_weight, in_proj_bias: input projection weight and bias.
+        static_k, static_v: static key and value used for attention operators.
+
+
+    Shape:
+        Inputs:
+        - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
+          the embedding dimension.
+        - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
+          the embedding dimension.
+        - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
+          the embedding dimension.
+        - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.
+          If a ByteTensor is provided, the non-zero positions will be ignored while the zero positions
+          will be unchanged. If a BoolTensor is provided, the positions with the
+          value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
+        - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.
+          3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,
+          S is the source sequence length. attn_mask ensures that position i is allowed to attend the unmasked
+          positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
+          while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
+          are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
+          is provided, it will be added to the attention weight.
+        - static_k: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,
+          N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.
+        - static_v: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,
+          N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.
+
+        Outputs:
+        - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
+          E is the embedding dimension.
+        - attn_output_weights: :math:`(N, L, S)` where N is the batch size,
+          L is the target sequence length, S is the source sequence length.
+    """
+    # tens_ops = (query, key, value, in_proj_weight, in_proj_bias, bias_k, bias_v, out_proj_weight, out_proj_bias)
+    # if torch.overrides.has_torch_function(tens_ops):
+    #     return torch.overrides.handle_torch_function(
+    #         multi_head_attention_forward,
+    #         tens_ops,
+    #         query,
+    #         key,
+    #         value,
+    #         embed_dim_to_check,
+    #         num_heads,
+    #         in_proj_weight,
+    #         in_proj_bias,
+    #         bias_k,
+    #         bias_v,
+    #         add_zero_attn,
+    #         dropout_p,
+    #         out_proj_weight,
+    #         out_proj_bias,
+    #         training=training,
+    #         key_padding_mask=key_padding_mask,
+    #         need_weights=need_weights,
+    #         attn_mask=attn_mask,
+    #         use_separate_proj_weight=use_separate_proj_weight,
+    #         q_proj_weight=q_proj_weight,
+    #         k_proj_weight=k_proj_weight,
+    #         v_proj_weight=v_proj_weight,
+    #         static_k=static_k,
+    #         static_v=static_v,
+    #         attn_bias=attn_bias,
+    #     )
+    tgt_len, bsz, embed_dim = query.size()
+    assert embed_dim == embed_dim_to_check
+    # allow MHA to have different sizes for the feature dimension
+    assert key.size(0) == value.size(0) and key.size(1) == value.size(1)
+
+    if isinstance(embed_dim, torch.Tensor):
+        # embed_dim can be a tensor when JIT tracing
+        head_dim = embed_dim.div(num_heads, rounding_mode='trunc')
+    else:
+        head_dim = embed_dim // num_heads
+    assert head_dim * num_heads == embed_dim, "embed_dim must be divisible by num_heads"
+    scaling = float(head_dim * scale_factor) ** -0.5
+
+    if not use_separate_proj_weight:
+        if (query is key or torch.equal(query, key)) and (key is value or torch.equal(key, value)):
+            # self-attention
+            q, k, v = linear(query, in_proj_weight, in_proj_bias).chunk(3, dim=-1)
+
+        elif key is value or torch.equal(key, value):
+            # encoder-decoder attention
+            # This is inline in_proj function with in_proj_weight and in_proj_bias
+            _b = in_proj_bias
+            _start = 0
+            _end = embed_dim
+            _w = in_proj_weight[_start:_end, :]
+            if _b is not None:
+                _b = _b[_start:_end]
+            q = linear(query, _w, _b)
+
+            if key is None:
+                assert value is None
+                k = None
+                v = None
+            else:
+
+                # This is inline in_proj function with in_proj_weight and in_proj_bias
+                _b = in_proj_bias
+                _start = embed_dim
+                _end = None
+                _w = in_proj_weight[_start:, :]
+                if _b is not None:
+                    _b = _b[_start:]
+                k, v = linear(key, _w, _b).chunk(2, dim=-1)
+
+        else:
+            # This is inline in_proj function with in_proj_weight and in_proj_bias
+            _b = in_proj_bias
+            _start = 0
+            _end = embed_dim
+            _w = in_proj_weight[_start:_end, :]
+            if _b is not None:
+                _b = _b[_start:_end]
+            q = linear(query, _w, _b)
+
+            # This is inline in_proj function with in_proj_weight and in_proj_bias
+            _b = in_proj_bias
+            _start = embed_dim
+            _end = embed_dim * 2
+            _w = in_proj_weight[_start:_end, :]
+            if _b is not None:
+                _b = _b[_start:_end]
+            k = linear(key, _w, _b)
+
+            # This is inline in_proj function with in_proj_weight and in_proj_bias
+            _b = in_proj_bias
+            _start = embed_dim * 2
+            _end = None
+            _w = in_proj_weight[_start:, :]
+            if _b is not None:
+                _b = _b[_start:]
+            v = linear(value, _w, _b)
+    else:
+        q_proj_weight_non_opt = torch.jit._unwrap_optional(q_proj_weight)
+        len1, len2 = q_proj_weight_non_opt.size()
+        assert len1 == embed_dim and len2 == query.size(-1)
+
+        k_proj_weight_non_opt = torch.jit._unwrap_optional(k_proj_weight)
+        len1, len2 = k_proj_weight_non_opt.size()
+        assert len1 == embed_dim and len2 == key.size(-1)
+
+        v_proj_weight_non_opt = torch.jit._unwrap_optional(v_proj_weight)
+        len1, len2 = v_proj_weight_non_opt.size()
+        assert len1 == embed_dim and len2 == value.size(-1)
+
+        if in_proj_bias is not None:
+            q = linear(query, q_proj_weight_non_opt, in_proj_bias[0:embed_dim])
+            k = linear(key, k_proj_weight_non_opt, in_proj_bias[embed_dim : (embed_dim * 2)])
+            v = linear(value, v_proj_weight_non_opt, in_proj_bias[(embed_dim * 2) :])
+        else:
+            q = linear(query, q_proj_weight_non_opt, in_proj_bias)
+            k = linear(key, k_proj_weight_non_opt, in_proj_bias)
+            v = linear(value, v_proj_weight_non_opt, in_proj_bias)
+    q = q * scaling
+
+    if attn_mask is not None:
+        assert (
+            attn_mask.dtype == torch.float32
+            or attn_mask.dtype == torch.float64
+            or attn_mask.dtype == torch.float16
+            or attn_mask.dtype == torch.uint8
+            or attn_mask.dtype == torch.bool
+        ), "Only float, byte, and bool types are supported for attn_mask, not {}".format(attn_mask.dtype)
+        if attn_mask.dtype == torch.uint8:
+            warnings.warn("Byte tensor for attn_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.")
+            attn_mask = attn_mask.to(torch.bool)
+
+        if attn_mask.dim() == 2:
+            attn_mask = attn_mask.unsqueeze(0)
+            if list(attn_mask.size()) != [1, query.size(0), key.size(0)]:
+                raise RuntimeError("The size of the 2D attn_mask is not correct.")
+        elif attn_mask.dim() == 3:
+            if list(attn_mask.size()) != [bsz * num_heads, query.size(0), key.size(0)]:
+                raise RuntimeError("The size of the 3D attn_mask is not correct.")
+        else:
+            raise RuntimeError("attn_mask's dimension {} is not supported".format(attn_mask.dim()))
+        # attn_mask's dim is 3 now.
+
+    # convert ByteTensor key_padding_mask to bool
+    if key_padding_mask is not None and key_padding_mask.dtype == torch.uint8:
+        warnings.warn(
+            "Byte tensor for key_padding_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead."
+        )
+        key_padding_mask = key_padding_mask.to(torch.bool)
+
+    if bias_k is not None and bias_v is not None:
+        if static_k is None and static_v is None:
+            k = torch.cat([k, bias_k.repeat(1, bsz, 1)])
+            v = torch.cat([v, bias_v.repeat(1, bsz, 1)])
+            if attn_mask is not None:
+                attn_mask = pad(attn_mask, (0, 1))
+            if key_padding_mask is not None:
+                key_padding_mask = pad(key_padding_mask, (0, 1))
+        else:
+            assert static_k is None, "bias cannot be added to static key."
+            assert static_v is None, "bias cannot be added to static value."
+    else:
+        assert bias_k is None
+        assert bias_v is None
+
+    q = q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
+    if k is not None:
+        k = k.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
+    if v is not None:
+        v = v.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
+
+    if static_k is not None:
+        assert static_k.size(0) == bsz * num_heads
+        assert static_k.size(2) == head_dim
+        k = static_k
+
+    if static_v is not None:
+        assert static_v.size(0) == bsz * num_heads
+        assert static_v.size(2) == head_dim
+        v = static_v
+
+    src_len = k.size(1)
+
+    if key_padding_mask is not None:
+        assert key_padding_mask.size(0) == bsz
+        assert key_padding_mask.size(1) == src_len
+
+    if add_zero_attn:
+        src_len += 1
+        k = torch.cat([k, torch.zeros((k.size(0), 1) + k.size()[2:], dtype=k.dtype, device=k.device)], dim=1)
+        v = torch.cat([v, torch.zeros((v.size(0), 1) + v.size()[2:], dtype=v.dtype, device=v.device)], dim=1)
+        if attn_mask is not None:
+            attn_mask = pad(attn_mask, (0, 1))
+        if key_padding_mask is not None:
+            key_padding_mask = pad(key_padding_mask, (0, 1))
+
+    attn_output_weights = torch.bmm(q, k.transpose(1, 2))
+    assert list(attn_output_weights.size()) == [bsz * num_heads, tgt_len, src_len]
+
+    if attn_bias is not None:
+        assert list(attn_bias.size()) == [bsz * num_heads, tgt_len, src_len]
+        attn_output_weights += attn_bias
+
+    if attn_mask is not None:
+        if attn_mask.dtype == torch.bool:
+            attn_output_weights.masked_fill_(attn_mask, float("-inf"))
+        else:
+            attn_output_weights += attn_mask
+
+    if key_padding_mask is not None:
+        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
+        attn_output_weights = attn_output_weights.masked_fill(
+            key_padding_mask.unsqueeze(1).unsqueeze(2),
+            float("-inf"),
+        )
+        attn_output_weights = attn_output_weights.view(bsz * num_heads, tgt_len, src_len)
+
+    attn_output_weights = softmax(attn_output_weights, dim=-1)
+    attn_output_weights = dropout(attn_output_weights, p=dropout_p, training=training)
+
+    attn_output = torch.bmm(attn_output_weights, v)
+    assert list(attn_output.size()) == [bsz * num_heads, tgt_len, head_dim]
+    attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
+    attn_output = linear(attn_output, out_proj_weight, out_proj_bias)
+
+    if need_weights:
+        # average attention weights over heads
+        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
+        return attn_output, attn_output_weights.sum(dim=1) / num_heads
+    else:
+        return attn_output, None
+
+
+class MultiheadAttention(Module):
+    r"""Allows the model to jointly attend to information
+    from different representation subspaces.
+    See `Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_
+
+    .. math::
+        \text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O
+
+    where :math:`head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)`.
+
+    Args:
+        embed_dim: total dimension of the model.
+        num_heads: parallel attention heads.
+        dropout: a Dropout layer on attn_output_weights. Default: 0.0.
+        bias: add bias as module parameter. Default: True.
+        add_bias_kv: add bias to the key and value sequences at dim=0.
+        add_zero_attn: add a new batch of zeros to the key and
+                       value sequences at dim=1.
+        kdim: total number of features in key. Default: None.
+        vdim: total number of features in value. Default: None.
+
+    Note that if :attr:`kdim` and :attr:`vdim` are None, they will be set
+    to :attr:`embed_dim` such that query, key, and value have the same
+    number of features.
+
+    Examples::
+
+        >>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
+        >>> attn_output, attn_output_weights = multihead_attn(query, key, value)
+    """
+    bias_k: Optional[torch.Tensor]
+    bias_v: Optional[torch.Tensor]
+
+    def __init__(self, embed_dim, num_heads, dropout=0., bias=True, add_bias_kv=False, scale_factor=1, add_zero_attn=False, kdim=None, vdim=None):
+        super(MultiheadAttention, self).__init__()
+        self.embed_dim = embed_dim
+        self.kdim = kdim if kdim is not None else embed_dim
+        self.vdim = vdim if vdim is not None else embed_dim
+        self._qkv_same_embed_dim = self.kdim == embed_dim and self.vdim == embed_dim
+        self.scale_factor = scale_factor
+
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.head_dim = embed_dim // num_heads
+        assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
+
+        if self._qkv_same_embed_dim is False:
+            self.q_proj_weight = Parameter(torch.empty(embed_dim, embed_dim))
+            self.k_proj_weight = Parameter(torch.empty(embed_dim, self.kdim))
+            self.v_proj_weight = Parameter(torch.empty(embed_dim, self.vdim))
+            self.register_parameter('in_proj_weight', None)
+        else:
+            self.in_proj_weight = Parameter(torch.empty(3 * embed_dim, embed_dim))
+            self.register_parameter('q_proj_weight', None)
+            self.register_parameter('k_proj_weight', None)
+            self.register_parameter('v_proj_weight', None)
+
+        if bias:
+            self.in_proj_bias = Parameter(torch.empty(3 * embed_dim))
+        else:
+            self.register_parameter('in_proj_bias', None)
+        self.out_proj = Linear(embed_dim, embed_dim, bias=bias)
+
+        if add_bias_kv:
+            self.bias_k = Parameter(torch.empty(1, 1, embed_dim))
+            self.bias_v = Parameter(torch.empty(1, 1, embed_dim))
+        else:
+            self.bias_k = self.bias_v = None
+
+        self.add_zero_attn = add_zero_attn
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        if self._qkv_same_embed_dim:
+            xavier_uniform_(self.in_proj_weight)
+        else:
+            xavier_uniform_(self.q_proj_weight)
+            xavier_uniform_(self.k_proj_weight)
+            xavier_uniform_(self.v_proj_weight)
+
+        if self.in_proj_bias is not None:
+            constant_(self.in_proj_bias, 0.)
+            constant_(self.out_proj.bias, 0.)
+        if self.bias_k is not None:
+            xavier_normal_(self.bias_k)
+        if self.bias_v is not None:
+            xavier_normal_(self.bias_v)
+
+    def __setstate__(self, state):
+        # Support loading old MultiheadAttention checkpoints generated by v1.1.0
+        if '_qkv_same_embed_dim' not in state:
+            state['_qkv_same_embed_dim'] = True
+
+        super(MultiheadAttention, self).__setstate__(state)
+
+    def forward(self, query: Tensor, key: Tensor, value: Tensor, key_padding_mask: Optional[Tensor] = None,
+                need_weights: bool = True, attn_mask: Optional[Tensor] = None, attn_bias: Optional[Tensor] = None) -> Tuple[Tensor, Optional[Tensor]]:
+        r"""
+    Args:
+        query, key, value: map a query and a set of key-value pairs to an output.
+            See "Attention Is All You Need" for more details.
+        key_padding_mask: if provided, specified padding elements in the key will
+            be ignored by the attention. When given a binary mask and a value is True,
+            the corresponding value on the attention layer will be ignored. When given
+            a byte mask and a value is non-zero, the corresponding value on the attention
+            layer will be ignored
+        need_weights: output attn_output_weights.
+        attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
+            the batches while a 3D mask allows to specify a different mask for the entries of each batch.
+
+    Shapes for inputs:
+        - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
+          the embedding dimension.
+        - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
+          the embedding dimension.
+        - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
+          the embedding dimension.
+        - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.
+          If a ByteTensor is provided, the non-zero positions will be ignored while the position
+          with the zero positions will be unchanged. If a BoolTensor is provided, the positions with the
+          value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
+        - attn_mask: if a 2D mask: :math:`(L, S)` where L is the target sequence length, S is the
+          source sequence length.
+
+          If a 3D mask: :math:`(N\cdot\text{num\_heads}, L, S)` where N is the batch size, L is the target sequence
+          length, S is the source sequence length. ``attn_mask`` ensure that position i is allowed to attend
+          the unmasked positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
+          while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
+          is not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
+          is provided, it will be added to the attention weight.
+
+    Shapes for outputs:
+        - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
+          E is the embedding dimension.
+        - attn_output_weights: :math:`(N, L, S)` where N is the batch size,
+          L is the target sequence length, S is the source sequence length.
+        """
+        if not self._qkv_same_embed_dim:
+            return multi_head_attention_forward(
+                query, key, value, self.embed_dim, self.num_heads,
+                self.in_proj_weight, self.in_proj_bias,
+                self.bias_k, self.bias_v, self.add_zero_attn,
+                self.dropout, self.out_proj.weight, self.out_proj.bias,
+                training=self.training,
+                key_padding_mask=key_padding_mask, need_weights=need_weights,
+                attn_mask=attn_mask, attn_bias=attn_bias,
+                use_separate_proj_weight=True,
+                q_proj_weight=self.q_proj_weight, k_proj_weight=self.k_proj_weight,
+                v_proj_weight=self.v_proj_weight)
+        else:
+            return multi_head_attention_forward(
+                query, key, value, self.embed_dim, self.num_heads,
+                self.in_proj_weight, self.in_proj_bias,
+                self.bias_k, self.bias_v, self.add_zero_attn,
+                self.dropout, self.out_proj.weight, self.out_proj.bias,
+                training=self.training,
+                key_padding_mask=key_padding_mask, need_weights=need_weights,
+                attn_mask=attn_mask, attn_bias=attn_bias)
+
+class TransformerEncoderLayer(Module):
+    r"""TransformerEncoderLayer is made up of self-attn and feedforward network.
+    This standard encoder layer is based on the paper "Attention Is All You Need".
+    Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
+    Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in
+    Neural Information Processing Systems, pages 6000-6010. Users may modify or implement
+    in a different way during application.
+
+    Args:
+        d_model: the number of expected features in the input (required).
+        nhead: the number of heads in the multiheadattention models (required).
+        dim_feedforward: the dimension of the feedforward network model (default=2048).
+        dropout: the dropout value (default=0.1).
+        activation: the activation function of intermediate layer, relu or gelu (default=relu).
+
+    Examples::
+        >>> encoder_layer = TransformerEncoderLayer(d_model=512, nhead=8)
+        >>> src = torch.rand(10, 32, 512)
+        >>> out = encoder_layer(src)
+    """
+
+    def __init__(self, d_model, nhead, attn_scale_factor=1, dim_feedforward=2048, dropout=0.1, activation="relu"):
+        super(TransformerEncoderLayer, self).__init__()
+        self.self_attn = MultiheadAttention(d_model, nhead, scale_factor=attn_scale_factor, dropout=dropout)
+        # Implementation of Feedforward model
+        self.linear1 = Linear(d_model, dim_feedforward)
+        self.dropout = Dropout(dropout)
+        self.linear2 = Linear(dim_feedforward, d_model)
+
+        self.norm1 = LayerNorm(d_model)
+        self.norm2 = LayerNorm(d_model)
+        self.dropout1 = Dropout(dropout)
+        self.dropout2 = Dropout(dropout)
+
+        self.activation = _get_activation_fn(activation)
+
+    def __setstate__(self, state):
+        if 'activation' not in state:
+            state['activation'] = F.relu
+        super(TransformerEncoderLayer, self).__setstate__(state)
+
+    def forward(self, src: Tensor, src_mask: Optional[Tensor] = None, attn_bias: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
+        r"""Pass the input through the encoder layer.
+
+        Args:
+            src: the sequence to the encoder layer (required).
+            src_mask: the mask for the src sequence (optional).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
+
+        Shape:
+            see the docs in Transformer class.
+        """
+        src2 = self.self_attn(src, src, src, attn_mask=src_mask, attn_bias=attn_bias,
+                              key_padding_mask=src_key_padding_mask)[0]
+        src = src + self.dropout1(src2)
+        src = self.norm1(src)
+        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
+        src = src + self.dropout2(src2)
+        src = self.norm2(src)
+        return src
+
+def _get_clones(module, N):
+    return ModuleList([copy.deepcopy(module) for i in range(N)])
+
+
+def _get_activation_fn(activation):
+    if activation == "relu":
+        return F.relu
+    elif activation == "gelu":
+        return F.gelu
+
+    raise RuntimeError("activation should be relu/gelu, not {}".format(activation))
+
+
+class TransformerEncoder(Module):
+    r"""TransformerEncoder is a stack of N encoder layers
+
+    Args:
+        encoder_layer: an instance of the TransformerEncoderLayer() class (required).
+        num_layers: the number of sub-encoder-layers in the encoder (required).
+        norm: the layer normalization component (optional).
+
+    Examples::
+        >>> encoder_layer = TransformerEncoderLayer(d_model=512, nhead=8)
+        >>> transformer_encoder = TransformerEncoder(encoder_layer, num_layers=6)
+        >>> src = torch.rand(10, 32, 512)
+        >>> out = transformer_encoder(src)
+    """
+    __constants__ = ['norm']
+
+    def __init__(self, encoder_layer, num_layers, norm=None):
+        super(TransformerEncoder, self).__init__()
+        self.layers = _get_clones(encoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+
+    def forward(self, src: Tensor, mask: Optional[Tensor] = None, attn_bias: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
+        r"""Pass the input through the encoder layers in turn.
+
+        Args:
+            src: the sequence to the encoder (required).
+            mask: the mask for the src sequence (optional).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
+
+        Shape:
+            see the docs in Transformer class.
+        """
+        output = src
+
+        for mod in self.layers:
+            output = mod(output, src_mask=mask, attn_bias=attn_bias, src_key_padding_mask=src_key_padding_mask)
+
+        if self.norm is not None:
+            output = self.norm(output)
+
+        return output

code/tabular_data_loaders.py

@@ -0,0 +1,676 @@
+# TabuLM — tabular dataset and data loading utilities
+# Extends KinyaBERT's data pipeline with:
+#   - TabularParsedToken: ParsedToken + (row_id, col_id, cell_type)
+#   - process_tabular_sentence: masking aware of cell boundaries
+#   - TabularKBCorpusDataset: loads CSV tables, no libkinlp.so required
+#   - tabular_collate_wrapper: extends morpho_seq_collate_wrapper
+#   - tabulm_model_forward: unpacks tabular batch for the model
+
+from __future__ import print_function, division
+
+import math
+import os
+import random
+import sys
+import warnings
+from typing import List, Optional, Tuple
+
+import numpy as np
+import torch
+import youtokentome as yttm
+from torch.utils.data import Dataset
+
+from kinyabert_utils import time_now
+from morpho_data_loaders import (
+    KBVocab, AffixSetVocab, ParsedToken, morpho_seq_collate_wrapper
+)
+from tabular_serializer import (
+    CellType, NUM_CELL_TYPES, TableCell,
+    serialize_csv, table_cells_to_text, WordMeta
+)
+from morpho_stub import parse_text_stub
+
+warnings.filterwarnings("ignore")
+
+
+# ── TabularParsedToken ────────────────────────────────────────────────────────
+
+class TabularParsedToken(ParsedToken):
+    """ParsedToken extended with table-grid coordinates."""
+
+    def __init__(self, row_id: int, col_id: int, cell_type: int, **kwargs):
+        super().__init__(**kwargs)
+        self.row_id = row_id
+        self.col_id = col_id
+        self.cell_type = cell_type
+
+
+# ── process_tabular_sentence ─────────────────────────────────────────────────
+
+def process_tabular_sentence(
+    args,
+    parsed_tokens_list: List[TabularParsedToken],
+    add_cls: bool,
+    kv: KBVocab,
+    affix_set_vocab: AffixSetVocab,
+    mcr_masked_cells: set,          # set of (row_id, col_id) selected for MCR masking
+    ctp_col_labels: dict,           # {col_id: cell_type_int} for CTP targets
+):
+    """
+    Extended version of process_parsed_sentence that additionally:
+      - Tracks row_id / col_id / cell_type per sequence position
+      - Applies MCR (cell-level) masking in addition to word-level masking
+      - Collects CTP prediction targets (column type labels)
+    """
+    pos_tags = []
+    stems = []
+    afsets = [] if args.use_afsets else None
+    affixes = []
+    tokens_lengths = []
+    row_ids = []
+    col_ids = []
+    cell_types = []
+
+    predicted_stems = []
+    predicted_afsets = [] if args.use_afsets else None
+    predicted_affixes = [] if args.predict_affixes else None
+    predicted_tokens_idx = []
+    predicted_tokens_affixes_idx = [] if args.predict_affixes else None
+    predicted_tokens_affixes_lengths = [] if args.predict_affixes else None
+
+    # MCR-specific bookkeeping
+    mcr_predicted_stems = []
+    mcr_predicted_tokens_idx = []
+
+    def _add_special(vocab_key: str, r: int = 0, c: int = 0, ct: int = 0):
+        pos_tags.append(kv.pos_tag_vocab[vocab_key])
+        stems.append(kv.reduced_stem_vocab[vocab_key])
+        if args.use_afsets:
+            afsets.append(affix_set_vocab.affix_set_to_idx(vocab_key))
+        tokens_lengths.append(0)
+        row_ids.append(r)
+        col_ids.append(c)
+        cell_types.append(ct)
+
+    if add_cls:
+        _add_special('<CLS>')
+
+    if not parsed_tokens_list:
+        _add_special('<SEP>')
+        return _pack_result(
+            pos_tags, stems, afsets, affixes, tokens_lengths,
+            row_ids, col_ids, cell_types,
+            predicted_stems, predicted_afsets, predicted_affixes,
+            predicted_tokens_idx, predicted_tokens_affixes_idx,
+            predicted_tokens_affixes_lengths,
+            mcr_predicted_stems, mcr_predicted_tokens_idx,
+        )
+
+    for ptoken in parsed_tokens_list:
+        is_mcr_masked = (ptoken.row_id, ptoken.col_id) in mcr_masked_cells
+        r, c, ct = ptoken.row_id, ptoken.col_id, ptoken.cell_type
+
+        for sidx in ptoken.stem_idx:
+            unchanged = True
+            predict_word = False
+            rval = random.random()
+
+            if is_mcr_masked:
+                # MCR: mask the entire cell unconditionally
+                pos_tags.append(kv.pos_tag_vocab['<MSK>'])
+                stems.append(kv.reduced_stem_vocab['<MSK>'])
+                if args.use_afsets:
+                    afsets.append(affix_set_vocab.affix_set_to_idx('<MSK>'))
+                tokens_lengths.append(0)
+                mcr_predicted_stems.append(kv.mapped_stem_vocab_idx[sidx])
+                mcr_predicted_tokens_idx.append(len(tokens_lengths) - 1)
+                unchanged = False
+
+            else:
+                # Standard word-level masking (15%)
+                if rval <= 0.15:
+                    predict_word = True
+                    rval /= 0.15
+                    if rval < 0.8:
+                        unchanged = False
+                        pos_tags.append(kv.pos_tag_vocab['<MSK>'])
+                        stems.append(kv.reduced_stem_vocab['<MSK>'])
+                        if args.use_afsets:
+                            afsets.append(affix_set_vocab.affix_set_to_idx('<MSK>'))
+                        if (rval / 0.8) < 0.3:
+                            affixes.extend(ptoken.affixes_idx)
+                            tokens_lengths.append(len(ptoken.affixes_idx))
+                        else:
+                            tokens_lengths.append(0)
+                    elif rval < 0.9:
+                        unchanged = False
+                        rnd_pos = random.randint(kv.pos_tag_vocab['<UNK>'], len(kv.pos_tag_vocab) - 1)
+                        rnd_stem = random.randint(kv.reduced_stem_vocab['<UNK>'], len(kv.reduced_stem_vocab) - 1)
+                        pos_tags.append(rnd_pos)
+                        stems.append(rnd_stem)
+                        if args.use_afsets:
+                            afsets.append(affix_set_vocab.random_idx())
+                        tokens_lengths.append(0)
+
+            if unchanged:
+                pos_tags.append(ptoken.pos_tag_idx)
+                stems.append(kv.mapped_stem_vocab_idx[sidx])
+                if args.use_afsets:
+                    afsets.append(affix_set_vocab.affix_set_to_idx(ptoken.affix_set_key()))
+                affixes.extend(ptoken.affixes_idx)
+                tokens_lengths.append(len(ptoken.affixes_idx))
+
+            row_ids.append(r)
+            col_ids.append(c)
+            cell_types.append(ct)
+
+            if predict_word:
+                predicted_stems.append(kv.mapped_stem_vocab_idx[sidx])
+                predicted_tokens_idx.append(len(tokens_lengths) - 1)
+                if args.use_afsets:
+                    predicted_afsets.append(affix_set_vocab.affix_set_to_idx(ptoken.affix_set_key()))
+                if args.predict_affixes:
+                    predicted_affixes.extend(ptoken.affixes_idx)
+                    if ptoken.affixes_idx:
+                        predicted_tokens_affixes_idx.append(len(predicted_tokens_idx) - 1)
+                        predicted_tokens_affixes_lengths.append(len(ptoken.affixes_idx))
+
+    return _pack_result(
+        pos_tags, stems, afsets, affixes, tokens_lengths,
+        row_ids, col_ids, cell_types,
+        predicted_stems, predicted_afsets, predicted_affixes,
+        predicted_tokens_idx, predicted_tokens_affixes_idx,
+        predicted_tokens_affixes_lengths,
+        mcr_predicted_stems, mcr_predicted_tokens_idx,
+    )
+
+
+def _pack_result(pos_tags, stems, afsets, affixes, tokens_lengths,
+                 row_ids, col_ids, cell_types,
+                 predicted_stems, predicted_afsets, predicted_affixes,
+                 predicted_tokens_idx, predicted_tokens_affixes_idx,
+                 predicted_tokens_affixes_lengths,
+                 mcr_predicted_stems, mcr_predicted_tokens_idx):
+    return (
+        pos_tags, stems, afsets, affixes, tokens_lengths,
+        row_ids, col_ids, cell_types,
+        predicted_stems, predicted_afsets, predicted_affixes,
+        predicted_tokens_idx, predicted_tokens_affixes_idx,
+        predicted_tokens_affixes_lengths,
+        mcr_predicted_stems, mcr_predicted_tokens_idx,
+    )
+
+
+# ── gather_tabular_itemized_data ─────────────────────────────────────────────
+
+def gather_tabular_itemized_data(
+    args,
+    table_cells_list: List[List[TableCell]],  # list of tables, each a list of cells
+    kb_vocab: KBVocab,
+    affix_set_vocab: AffixSetVocab,
+    bpe: yttm.BPE,
+    max_seq_len: int,
+    max_batch_items: int,
+    mcr_cell_mask_rate: float = 0.15,
+    ctp_mask_rate: float = 0.50,
+):
+    """
+    Builds a list of itemized training data items from a list of tables.
+    Each item is a tuple compatible with tabular_collate_wrapper.
+    """
+    itemized_data = []
+    table_idx = 0
+
+    seq_pos_tags, seq_stems, seq_afsets, seq_affixes, seq_tokens_lengths = [], [], [], [], []
+    seq_row_ids, seq_col_ids, seq_cell_types = [], [], []
+    seq_predicted_stems = []
+    seq_predicted_afsets = [] if args.use_afsets else None
+    seq_predicted_affixes = [] if args.predict_affixes else None
+    seq_predicted_tokens_idx = []
+    seq_predicted_tokens_affixes_idx = [] if args.predict_affixes else None
+    seq_predicted_tokens_affixes_lengths = [] if args.predict_affixes else None
+    seq_mcr_predicted_stems, seq_mcr_predicted_tokens_idx = [], []
+    seq_ctp_labels: List[Tuple[int, int]] = []   # (position_in_sequence, cell_type_label)
+    seq_rrp_pair: Optional[Tuple[int, int, int]] = None  # (rowA_start, rowB_start, label)
+
+    add_cls = True
+    random.shuffle(table_cells_list)
+
+    for cells in table_cells_list:
+        if not cells:
+            continue
+
+        text, word_meta = table_cells_to_text(cells)
+        parsed_tokens_raw = parse_text_stub(text, kb_vocab, bpe)
+
+        if len(parsed_tokens_raw) != len(word_meta):
+            # Length mismatch can happen with empty cells — skip
+            continue
+
+        # Attach table coordinates to each ParsedToken
+        tabular_tokens: List[TabularParsedToken] = []
+        for pt, (r, c, ct) in zip(parsed_tokens_raw, word_meta):
+            tp = TabularParsedToken(
+                row_id=r, col_id=c, cell_type=ct,
+                surface_form=pt.surface_form,
+                decode_prob=pt.decode_prob,
+                tf_idf=pt.tf_idf,
+                pos_tag_id=pt.pos_tag_idx,
+                stem_ids=pt.stem_idx,
+            )
+            tp.morpho_slots_idx = pt.morpho_slots_idx
+            tp.affixes_idx = pt.affixes_idx
+            tabular_tokens.append(tp)
+
+        # Choose MCR cells to mask (15% of unique cells by (row_id, col_id))
+        unique_cells = set((t.row_id, t.col_id) for t in tabular_tokens if t.row_id > 0)
+        num_mcr = max(1, int(mcr_cell_mask_rate * len(unique_cells)))
+        mcr_masked_cells = set(random.sample(sorted(unique_cells), min(num_mcr, len(unique_cells))))
+
+        # CTP: determine column type labels (majority vote per column)
+        col_type_votes: dict = {}
+        for t in tabular_tokens:
+            if t.row_id > 1 and t.col_id > 0 and t.cell_type != int(CellType.PAD):
+                col_type_votes.setdefault(t.col_id, []).append(t.cell_type)
+
+        ctp_labels_for_table = {}
+        for col_id, votes in col_type_votes.items():
+            from collections import Counter
+            majority = Counter(votes).most_common(1)[0][0]
+            ctp_labels_for_table[col_id] = majority
+
+        (pos_tags, stems, afsets, affixes, tokens_lengths,
+         row_ids, col_ids, cell_types,
+         predicted_stems, predicted_afsets, predicted_affixes,
+         predicted_tokens_idx, predicted_tokens_affixes_idx,
+         predicted_tokens_affixes_lengths,
+         mcr_predicted_stems, mcr_predicted_tokens_idx) = process_tabular_sentence(
+            args, tabular_tokens, add_cls, kb_vocab, affix_set_vocab,
+            mcr_masked_cells, ctp_labels_for_table,
+        )
+        add_cls = False
+
+        if (len(seq_tokens_lengths) + len(tokens_lengths)) > max_seq_len:
+            # Flush existing buffer (only if non-empty)
+            if seq_tokens_lengths:
+                item = _pack_tabular_item(
+                    max_seq_len,
+                    seq_pos_tags, seq_stems, seq_afsets, seq_affixes, seq_tokens_lengths,
+                    seq_row_ids, seq_col_ids, seq_cell_types,
+                    seq_predicted_stems, seq_predicted_afsets, seq_predicted_affixes,
+                    seq_predicted_tokens_idx, seq_predicted_tokens_affixes_idx,
+                    seq_predicted_tokens_affixes_lengths,
+                    seq_mcr_predicted_stems, seq_mcr_predicted_tokens_idx,
+                    seq_ctp_labels,
+                )
+                itemized_data.append(item)
+                if len(itemized_data) >= max_batch_items:
+                    return itemized_data
+
+            # Reset buffer then fall through to add current tokens below
+            seq_pos_tags, seq_stems, seq_afsets, seq_affixes, seq_tokens_lengths = [], [], [], [], []
+            seq_row_ids, seq_col_ids, seq_cell_types = [], [], []
+            seq_predicted_stems = []
+            seq_predicted_afsets = [] if args.use_afsets else None
+            seq_predicted_affixes = [] if args.predict_affixes else None
+            seq_predicted_tokens_idx = []
+            seq_predicted_tokens_affixes_idx = [] if args.predict_affixes else None
+            seq_predicted_tokens_affixes_lengths = [] if args.predict_affixes else None
+            seq_mcr_predicted_stems, seq_mcr_predicted_tokens_idx = [], []
+            seq_ctp_labels = []
+            add_cls = True
+            # Truncate current table to max_seq_len before adding to fresh buffer
+            tokens_lengths = tokens_lengths[:max_seq_len]
+            pos_tags = pos_tags[:max_seq_len]
+            stems = stems[:max_seq_len]
+            row_ids = row_ids[:max_seq_len]
+            col_ids = col_ids[:max_seq_len]
+            cell_types = cell_types[:max_seq_len]
+            if afsets is not None:
+                afsets = afsets[:max_seq_len]
+            # Re-compute affix flat list up to the truncated tokens
+            affix_end = sum(tokens_lengths)
+            affixes = affixes[:affix_end]
+            # Trim prediction indices that are out of range
+            predicted_tokens_idx = [i for i in predicted_tokens_idx if i < max_seq_len]
+            predicted_stems = predicted_stems[:len(predicted_tokens_idx)]
+            if predicted_afsets is not None:
+                predicted_afsets = predicted_afsets[:len(predicted_tokens_idx)]
+            if predicted_tokens_affixes_idx is not None:
+                predicted_tokens_affixes_idx = [i for i in predicted_tokens_affixes_idx if i < len(predicted_tokens_idx)]
+                predicted_tokens_affixes_lengths = predicted_tokens_affixes_lengths[:len(predicted_tokens_affixes_idx)] if predicted_tokens_affixes_lengths else []
+                predicted_affixes = predicted_affixes[:sum(predicted_tokens_affixes_lengths)] if predicted_affixes else []
+            mcr_predicted_tokens_idx = [i for i in mcr_predicted_tokens_idx if i < max_seq_len]
+            mcr_predicted_stems = mcr_predicted_stems[:len(mcr_predicted_tokens_idx)]
+
+        offset = len(seq_predicted_tokens_idx)
+        if args.predict_affixes and predicted_tokens_affixes_idx:
+            seq_predicted_tokens_affixes_idx.extend(
+                [offset + i for i in predicted_tokens_affixes_idx]
+            )
+        seq_predicted_tokens_idx.extend(
+            [len(seq_tokens_lengths) + i for i in predicted_tokens_idx]
+        )
+
+        # CTP: record (absolute_position_of_first_header_token_in_col, label)
+        header_positions = {}
+        for abs_i, (r, c, ct_int) in enumerate(zip(row_ids, col_ids, cell_types)):
+            if ct_int == int(CellType.HEADER) and c not in header_positions:
+                header_positions[c] = len(seq_tokens_lengths) + abs_i
+        for col_id, label in ctp_labels_for_table.items():
+            if random.random() < ctp_mask_rate and col_id in header_positions:
+                seq_ctp_labels.append((header_positions[col_id], label))
+
+        seq_pos_tags.extend(pos_tags)
+        seq_stems.extend(stems)
+        if args.use_afsets and afsets is not None:
+            seq_afsets.extend(afsets)
+        seq_affixes.extend(affixes)
+        seq_tokens_lengths.extend(tokens_lengths)
+        seq_row_ids.extend(row_ids)
+        seq_col_ids.extend(col_ids)
+        seq_cell_types.extend(cell_types)
+        seq_predicted_stems.extend(predicted_stems)
+        if args.use_afsets and predicted_afsets is not None:
+            seq_predicted_afsets.extend(predicted_afsets)
+        if args.predict_affixes and predicted_affixes is not None:
+            seq_predicted_affixes.extend(predicted_affixes)
+            if predicted_tokens_affixes_lengths:
+                seq_predicted_tokens_affixes_lengths.extend(predicted_tokens_affixes_lengths)
+        seq_mcr_predicted_stems.extend(mcr_predicted_stems)
+        seq_mcr_predicted_tokens_idx.extend(
+            [len(seq_tokens_lengths) - len(tokens_lengths) + i for i in mcr_predicted_tokens_idx]
+        )
+
+    return itemized_data
+
+
+def _pack_tabular_item(
+    max_seq_len,
+    pos_tags, stems, afsets, affixes, tokens_lengths,
+    row_ids, col_ids, cell_types,
+    predicted_stems, predicted_afsets, predicted_affixes,
+    predicted_tokens_idx, predicted_tokens_affixes_idx,
+    predicted_tokens_affixes_lengths,
+    mcr_predicted_stems, mcr_predicted_tokens_idx,
+    ctp_labels,
+):
+    return (
+        max_seq_len,
+        # Original KinyaBERT fields
+        None,           # rel_pos_arr (unused for tabular)
+        pos_tags, stems, afsets, affixes, tokens_lengths,
+        predicted_stems, predicted_afsets, predicted_affixes,
+        predicted_tokens_idx, predicted_tokens_affixes_idx,
+        predicted_tokens_affixes_lengths,
+        # Tabular-specific fields
+        row_ids, col_ids, cell_types,
+        mcr_predicted_stems, mcr_predicted_tokens_idx,
+        ctp_labels,
+    )
+
+
+# ── TabularKBCorpusDataset ────────────────────────────────────────────────────
+
+class TabularKBCorpusDataset(Dataset):
+    """
+    Dataset that loads CSV tables from a directory and builds
+    tabular training items. Compatible with tabular_collate_wrapper.
+    """
+
+    def __init__(
+        self,
+        args,
+        kb_vocab: KBVocab,
+        affix_set_vocab: AffixSetVocab,
+        bpe_encoder: yttm.BPE,
+        csv_dir: str,
+        max_batch_items: int,
+        max_seq_len: int = 512,
+        max_rows: int = 64,
+        max_cols: int = 24,
+        rank: int = 0,
+    ):
+        self.max_seq_len = max_seq_len
+        self.max_batch_items = max_batch_items
+
+        # Discover CSV files
+        csv_files = [
+            os.path.join(csv_dir, f)
+            for f in os.listdir(csv_dir)
+            if f.lower().endswith('.csv')
+        ]
+        if not csv_files:
+            raise FileNotFoundError(f'No CSV files found in {csv_dir}')
+
+        if rank == 0:
+            print(time_now(), f'Loading {len(csv_files)} CSV tables from {csv_dir}')
+
+        # Load all tables
+        all_tables: List[List[TableCell]] = []
+        for fp in csv_files:
+            cells = serialize_csv(fp, max_rows=max_rows, max_cols=max_cols)
+            if cells:
+                all_tables.append(cells)
+
+        if rank == 0:
+            print(time_now(), f'{len(all_tables)} tables loaded successfully')
+
+        self.itemized_data = gather_tabular_itemized_data(
+            args, all_tables, kb_vocab, affix_set_vocab, bpe_encoder,
+            max_seq_len=max_seq_len,
+            max_batch_items=max_batch_items,
+        )
+
+        if rank == 0:
+            print(time_now(), f'{len(self.itemized_data)} training items prepared')
+
+    def __len__(self):
+        return len(self.itemized_data)
+
+    def __getitem__(self, idx):
+        return self.itemized_data[idx]
+
+
+# ── tabular_collate_wrapper ───────────────────────────────────────────────────
+
+def tabular_collate_wrapper(batch_items):
+    """
+    Collate function for TabularKBCorpusDataset.
+    Extends morpho_seq_collate_wrapper with tabular fields.
+    """
+    batch_input_sequence_lengths = []
+    batch_pos_tags, batch_stems, batch_afsets = [], [], []
+    batch_affixes, batch_tokens_lengths = [], []
+    batch_predicted_stems, batch_predicted_afsets = [], []
+    batch_predicted_affixes = []
+    batch_predicted_tokens_idx = []
+    batch_predicted_tokens_affixes_idx = []
+    batch_predicted_tokens_affixes_lengths = []
+
+    batch_row_ids, batch_col_ids, batch_cell_types = [], [], []
+    batch_mcr_predicted_stems, batch_mcr_predicted_tokens_idx = [], []
+    batch_ctp_labels: List[Tuple[int, int]] = []
+
+    for bidx, item in enumerate(batch_items):
+        (max_seq_len, _rel_pos_arr,
+         pos_tags, stems, afsets, affixes, tokens_lengths,
+         predicted_stems, predicted_afsets, predicted_affixes,
+         predicted_tokens_idx, predicted_tokens_affixes_idx,
+         predicted_tokens_affixes_lengths,
+         row_ids, col_ids, cell_types,
+         mcr_predicted_stems, mcr_predicted_tokens_idx,
+         ctp_labels) = item
+
+        # Offset affix-prediction indices
+        if predicted_tokens_affixes_idx is not None:
+            batch_predicted_tokens_affixes_idx.extend(
+                [(len(batch_predicted_tokens_idx) + t) for t in predicted_tokens_affixes_idx]
+            )
+
+        batch_predicted_tokens_idx.extend(
+            [(t, len(batch_input_sequence_lengths)) for t in predicted_tokens_idx]
+        )
+
+        # Offset MCR indices the same way
+        batch_mcr_predicted_tokens_idx.extend(
+            [(t, len(batch_input_sequence_lengths)) for t in mcr_predicted_tokens_idx]
+        )
+
+        # CTP labels: offset the absolute position
+        seq_start = sum(batch_input_sequence_lengths)  # not directly, but captured below
+        batch_ctp_labels.extend(ctp_labels)
+
+        batch_pos_tags.extend(pos_tags)
+        batch_stems.extend(stems)
+        if afsets is not None:
+            batch_afsets.extend(afsets)
+        batch_affixes.extend(affixes)
+        batch_tokens_lengths.extend(tokens_lengths)
+        batch_predicted_stems.extend(predicted_stems)
+        if predicted_afsets is not None:
+            batch_predicted_afsets.extend(predicted_afsets)
+        if predicted_affixes is not None:
+            batch_predicted_affixes.extend(predicted_affixes)
+        if predicted_tokens_affixes_lengths is not None:
+            batch_predicted_tokens_affixes_lengths.extend(predicted_tokens_affixes_lengths)
+
+        batch_row_ids.extend(row_ids)
+        batch_col_ids.extend(col_ids)
+        batch_cell_types.extend(cell_types)
+        batch_mcr_predicted_stems.extend(mcr_predicted_stems)
+        batch_input_sequence_lengths.append(len(tokens_lengths))
+
+    return (
+        batch_input_sequence_lengths,
+        None,  # rel_pos_arr unused
+        batch_pos_tags, batch_stems, batch_afsets,
+        batch_affixes, batch_tokens_lengths,
+        batch_predicted_stems, batch_predicted_afsets, batch_predicted_affixes,
+        batch_predicted_tokens_idx,
+        batch_predicted_tokens_affixes_idx,
+        batch_predicted_tokens_affixes_lengths,
+        # Tabular additions
+        batch_row_ids, batch_col_ids, batch_cell_types,
+        batch_mcr_predicted_stems, batch_mcr_predicted_tokens_idx,
+        batch_ctp_labels,
+    )
+
+
+# ── tabulm_model_forward ──────────────────────────────────────────────────────
+
+def tabulm_model_forward(args, data_item, model, device, tot_num_affixes):
+    """
+    Unpacks a tabular data_item and runs one forward pass through TabuLM.
+    Returns (total_loss, stem_loss, afset_loss, affix_loss, mcr_loss, ctp_loss, rrp_loss).
+    """
+    (batch_input_sequence_lengths,
+     _rel_pos_arr,
+     batch_pos_tags, batch_stems, batch_afsets,
+     batch_affixes, batch_tokens_lengths,
+     batch_predicted_stems, batch_predicted_afsets, batch_predicted_affixes,
+     batch_predicted_tokens_idx,
+     batch_predicted_tokens_affixes_idx,
+     batch_predicted_tokens_affixes_lengths,
+     batch_row_ids, batch_col_ids, batch_cell_types,
+     batch_mcr_predicted_stems, batch_mcr_predicted_tokens_idx,
+     batch_ctp_labels) = data_item
+
+    pos_tags = torch.tensor(batch_pos_tags, dtype=torch.long).to(device)
+    stems = torch.tensor(batch_stems, dtype=torch.long).to(device)
+    afsets = torch.tensor(batch_afsets, dtype=torch.long).to(device) if args.use_afsets and batch_afsets else None
+    affixes = torch.tensor(batch_affixes, dtype=torch.long).to(device)
+    row_ids = torch.tensor(batch_row_ids, dtype=torch.long).to(device)
+    col_ids = torch.tensor(batch_col_ids, dtype=torch.long).to(device)
+    cell_types = torch.tensor(batch_cell_types, dtype=torch.long).to(device)
+
+    max_seq = max(batch_input_sequence_lengths)
+    predicted_tokens_idx = torch.tensor(
+        [s * max_seq + t for t, s in batch_predicted_tokens_idx], dtype=torch.long
+    ).to(device)
+
+    predicted_tokens_affixes_idx = (
+        torch.tensor(batch_predicted_tokens_affixes_idx, dtype=torch.long).to(device)
+        if args.predict_affixes and batch_predicted_tokens_affixes_idx
+        else None
+    )
+
+    predicted_affixes_prob = None
+    if args.predict_affixes and batch_predicted_affixes and predicted_tokens_affixes_idx is not None:
+        from itertools import accumulate
+        lengths = batch_predicted_tokens_affixes_lengths
+        affix_groups = []
+        prev = 0
+        for ln in lengths:
+            sub = batch_predicted_affixes[prev: prev + ln]
+            vec = torch.zeros(tot_num_affixes)
+            for idx in sub:
+                if 0 < idx < tot_num_affixes:
+                    vec[idx] = 1.0
+            affix_groups.append(vec)
+            prev += ln
+        if affix_groups:
+            predicted_affixes_prob = torch.stack(affix_groups).to(device)
+            predicted_affixes_prob = predicted_affixes_prob / (predicted_affixes_prob.sum(1, keepdim=True).clamp(min=1.0))
+
+    predicted_stems = torch.tensor(batch_predicted_stems, dtype=torch.long).to(device)
+    predicted_afsets = (
+        torch.tensor(batch_predicted_afsets, dtype=torch.long).to(device)
+        if args.use_afsets and batch_predicted_afsets else None
+    )
+
+    # MCR indices and targets
+    mcr_tokens_idx = (
+        torch.tensor([s * max_seq + t for t, s in batch_mcr_predicted_tokens_idx], dtype=torch.long).to(device)
+        if batch_mcr_predicted_tokens_idx else None
+    )
+    mcr_stems = (
+        torch.tensor(batch_mcr_predicted_stems, dtype=torch.long).to(device)
+        if batch_mcr_predicted_stems else None
+    )
+
+    # CTP indices and labels
+    ctp_positions = (
+        torch.tensor([pos for pos, _ in batch_ctp_labels], dtype=torch.long).to(device)
+        if batch_ctp_labels else None
+    )
+    ctp_labels_t = (
+        torch.tensor([lbl for _, lbl in batch_ctp_labels], dtype=torch.long).to(device)
+        if batch_ctp_labels else None
+    )
+
+    # Ablation flags: disable objectives or structural embeddings
+    if getattr(args, 'no_mcr', False):
+        mcr_tokens_idx = None
+        mcr_stems = None
+    if getattr(args, 'no_ctp', False):
+        ctp_positions = None
+        ctp_labels_t = None
+    if getattr(args, 'no_tabular_emb', False):
+        # padding_idx=0 → embedding lookup returns zeros, so zeroing IDs kills all structural signal
+        row_ids = torch.zeros_like(row_ids)
+        col_ids = torch.zeros_like(col_ids)
+        cell_types = torch.zeros_like(cell_types)
+
+    return model(
+        args,
+        rel_pos_arr=None,
+        tokens_lengths=batch_tokens_lengths,
+        input_sequence_lengths=batch_input_sequence_lengths,
+        pos_tags=pos_tags,
+        stems=stems,
+        afsets=afsets,
+        affixes=affixes,
+        row_ids=row_ids,
+        col_ids=col_ids,
+        cell_types=cell_types,
+        predicted_tokens_idx=predicted_tokens_idx,
+        predicted_tokens_affixes_idx=predicted_tokens_affixes_idx,
+        predicted_stems=predicted_stems,
+        predicted_afsets=predicted_afsets,
+        predicted_affixes_prob=predicted_affixes_prob,
+        mcr_tokens_idx=mcr_tokens_idx,
+        mcr_stems=mcr_stems,
+        ctp_positions=ctp_positions,
+        ctp_labels=ctp_labels_t,
+    )

code/tabular_serializer.py

@@ -0,0 +1,175 @@
+# TabuLM — tabular data serializer
+# Converts CSV / dict-of-records tables into a flat token string
+# plus a parallel word_metadata list for row/column tracking.
+
+import csv
+import re
+from enum import IntEnum
+from typing import List, Tuple
+
+# ── Cell type taxonomy ────────────────────────────────────────────────────────
+
+class CellType(IntEnum):
+    PAD         = 0   # padding / special tokens with no cell identity
+    HEADER      = 1   # column header row
+    NUMERIC     = 2   # quantity, count, percentage, measurement
+    TEXT        = 3   # free-form text longer than a label
+    CATEGORICAL = 4   # short label / enum value
+    DATE        = 5   # year or full date
+
+NUM_CELL_TYPES = 6
+
+# ── Regex heuristics ──────────────────────────────────────────────────────────
+
+_NUMERIC_RE = re.compile(
+    r'^[\d,.\s]+(%|Frw|RWF|km|kg|ha|m|m²|L|MW|USD|acres|ha)?$',
+    re.IGNORECASE,
+)
+_DATE_RE = re.compile(
+    r'^(1[89]\d{2}|2[012]\d{2})(-\d{2}(-\d{2})?)?$|^\d{1,2}/\d{1,2}/\d{2,4}$'
+)
+_NULL_RE = re.compile(r'^[-–]$|^n/?a$|^null$|^none$|^$', re.IGNORECASE)
+
+
+# ── Core data structures ──────────────────────────────────────────────────────
+
+class TableCell:
+    """One cell in a table, with its grid coordinates and semantic type."""
+    __slots__ = ('content', 'row_id', 'col_id', 'cell_type')
+
+    def __init__(self, content: str, row_id: int, col_id: int, cell_type: CellType):
+        self.content = content.strip()
+        self.row_id = row_id     # 1-based; 1 = header row
+        self.col_id = col_id     # 1-based
+        self.cell_type = cell_type
+
+    def __repr__(self):
+        return (f'TableCell(r={self.row_id}, c={self.col_id}, '
+                f'type={self.cell_type.name}, "{self.content[:20]}")')
+
+
+# ── Cell type detection ───────────────────────────────────────────────────────
+
+def detect_cell_type(value: str, is_header: bool = False) -> CellType:
+    if is_header:
+        return CellType.HEADER
+    v = value.strip()
+    if _NULL_RE.match(v):
+        return CellType.TEXT
+    if _DATE_RE.match(v):
+        return CellType.DATE
+    if _NUMERIC_RE.match(v):
+        return CellType.NUMERIC
+    if len(v) <= 40 and '\n' not in v and ' ' not in v:
+        return CellType.CATEGORICAL
+    return CellType.TEXT
+
+
+# ── Table loading ─────────────────────────────────────────────────────────────
+
+def serialize_csv(filepath: str,
+                  max_rows: int = 64,
+                  max_cols: int = 24) -> List[TableCell]:
+    """Read a CSV file and return an ordered list of TableCell objects."""
+    cells: List[TableCell] = []
+    try:
+        with open(filepath, newline='', encoding='utf-8-sig') as f:
+            rows = [r for r in csv.reader(f) if any(c.strip() for c in r)]
+    except Exception:
+        return cells
+
+    if not rows:
+        return cells
+
+    header = rows[0][:max_cols]
+    for col_id, h in enumerate(header, start=1):
+        cells.append(TableCell(
+            h or f'col_{col_id}', row_id=1, col_id=col_id,
+            cell_type=CellType.HEADER,
+        ))
+
+    for row_offset, row in enumerate(rows[1: max_rows + 1], start=2):
+        for col_id, val in enumerate(row[:max_cols], start=1):
+            cells.append(TableCell(
+                val, row_id=row_offset, col_id=col_id,
+                cell_type=detect_cell_type(val),
+            ))
+
+    return cells
+
+
+def serialize_records(records: List[dict],
+                      max_rows: int = 64,
+                      max_cols: int = 24) -> List[TableCell]:
+    """Convert a list of dicts (e.g. from pandas .to_dict('records')) to cells."""
+    if not records:
+        return []
+    cells: List[TableCell] = []
+    keys = list(records[0].keys())[:max_cols]
+
+    for col_id, k in enumerate(keys, start=1):
+        cells.append(TableCell(
+            str(k), row_id=1, col_id=col_id,
+            cell_type=CellType.HEADER,
+        ))
+
+    for row_offset, rec in enumerate(records[:max_rows], start=2):
+        for col_id, k in enumerate(keys, start=1):
+            val = str(rec.get(k, ''))
+            cells.append(TableCell(
+                val, row_id=row_offset, col_id=col_id,
+                cell_type=detect_cell_type(val),
+            ))
+
+    return cells
+
+
+# ── Serialization ────────────────────────────────────────────────────────────���
+
+# WordMeta: (row_id, col_id, cell_type_int)
+WordMeta = Tuple[int, int, int]
+
+
+def table_cells_to_text(
+    cells: List[TableCell],
+) -> Tuple[str, List[WordMeta]]:
+    """
+    Flatten table cells into a single space-separated string with structure
+    tokens, plus a parallel per-token metadata list.
+
+    Special tokens emitted:
+      [TAB]  — start of header row
+      [ROW]  — start of any data row
+      [CEL]  — start of each individual cell
+
+    Returns:
+      text         — string ready for morpho_stub.parse_text_stub()
+      word_meta    — list of (row_id, col_id, cell_type) with one entry per
+                     space-separated token in `text` (including special tokens,
+                     which get row_id=0, col_id=0, cell_type=PAD)
+    """
+    parts: List[str] = []
+    word_meta: List[WordMeta] = []
+
+    PAD = (0, 0, int(CellType.PAD))
+    cur_row = -1
+
+    for cell in cells:
+        if cell.row_id != cur_row:
+            sep = '[TAB]' if cell.row_id == 1 else '[ROW]'
+            parts.append(sep)
+            word_meta.append(PAD)
+            cur_row = cell.row_id
+
+        parts.append('[CEL]')
+        word_meta.append((cell.row_id, cell.col_id, int(cell.cell_type)))
+
+        content_words = cell.content.split() if cell.content else ['[EMPTY]']
+        if not content_words:
+            content_words = ['[EMPTY]']
+
+        for w in content_words:
+            parts.append(w)
+            word_meta.append((cell.row_id, cell.col_id, int(cell.cell_type)))
+
+    return ' '.join(parts), word_meta

code/tabulm_model.py

@@ -0,0 +1,468 @@
+# TabuLM — Table-Aware Morphologically-Rich Language Model
+# Extends KinyaBERT_MorphoEncoder with:
+#   1. Row / column / cell-type embeddings (additive, same dim as seq_tr_d_model)
+#   2. Table structure attention bias (same-row, same-col, header bias per head)
+#   3. MCR head  — Masked Cell Recovery (reuses stem decoder)
+#   4. CTP head  — Column Type Prediction (4-class linear)
+#
+# Architecture change summary vs KinyaBERT:
+#   seq_tr_d_model   unchanged (768 with default args)
+#   attn_bias        = position_bias + table_structure_bias  (injected at morpho_transformer.py:333)
+#   input to seq-tr  = morpho_output + stem_embed + row_embed + col_embed + cell_type_embed  (additive)
+#   new losses       = MCR (NLL) + CTP (NLL)
+
+from __future__ import print_function, division
+
+import math
+import warnings
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.cuda.amp import custom_fwd
+from torch.nn.utils.rnn import pad_sequence
+
+from morpho_transformer import TransformerEncoder, TransformerEncoderLayer, MultiheadAttention
+from morpho_model import (
+    gelu, init_bert_params, BertLayerNorm,
+    KinyaBERT_MorphoEncoder,
+)
+from tabular_serializer import NUM_CELL_TYPES
+
+warnings.filterwarnings("ignore")
+
+# ── Maximum table dimensions supported ───────────────────────────────────────
+MAX_ROWS = 128
+MAX_COLS = 64
+
+
+# ── Column type label set for CTP ────────────────────────────────────────────
+# We predict: NUMERIC(0) / TEXT(1) / CATEGORICAL(2) / DATE(3)
+# (HEADER is the header row itself, not a prediction target)
+CTP_NUM_LABELS = 4
+CTP_TYPE_MAP = {2: 0, 3: 1, 4: 2, 5: 3}   # CellType int → CTP label index
+
+
+class TabuLM_Encoder(KinyaBERT_MorphoEncoder):
+    """
+    KinyaBERT_MorphoEncoder extended with tabular structure awareness.
+
+    Three new components are added to the sequence transformer tier:
+
+    (a) Additive embeddings: row_embedding, col_embedding, cell_type_embedding
+        — same dimension as seq_tr_d_model, added element-wise to each word
+          vector before it enters the sequence transformer (same way BERT adds
+          absolute position embeddings).
+
+    (b) Table structure attention bias: a per-head learned scalar that is added
+        to the pre-softmax attention weights whenever two tokens share the same
+        row, same column, or when either token is in the header row.
+        Injected at morpho_transformer.py line 333 on top of position bias.
+
+    (c) MCR prediction head: reuses the inherited stem decoder to score
+        cell-level masked tokens (no new parameters needed).
+
+    (d) CTP prediction head: a 2-layer linear that maps from seq_tr_d_model
+        to CTP_NUM_LABELS classes for column type prediction.
+    """
+
+    def __init__(self, *args,
+                 max_rows: int = MAX_ROWS,
+                 max_cols: int = MAX_COLS,
+                 **kwargs):
+        super().__init__(*args, **kwargs)
+
+        d = self.seq_tr_d_model
+
+        # (a) Tabular embeddings (additive, padding_idx=0 is the PAD/special token slot)
+        self.row_embedding = nn.Embedding(max_rows + 2, d, padding_idx=0)
+        self.col_embedding = nn.Embedding(max_cols + 2, d, padding_idx=0)
+        self.cell_type_embedding = nn.Embedding(NUM_CELL_TYPES + 1, d, padding_idx=0)
+
+        # (b) Table structure bias: one scalar per attention head, for each relation
+        self.row_attn_bias = nn.Parameter(torch.zeros(self.seq_tr_nhead))
+        self.col_attn_bias = nn.Parameter(torch.zeros(self.seq_tr_nhead))
+        self.header_attn_bias = nn.Parameter(torch.zeros(self.seq_tr_nhead))
+
+        # (d) CTP head
+        self.ctp_head = nn.Sequential(
+            nn.Linear(d, d // 2),
+            nn.GELU(),
+            BertLayerNorm(d // 2, eps=1e-6),
+            nn.Linear(d // 2, CTP_NUM_LABELS),
+        )
+
+        self.apply(init_bert_params)
+
+    # ── Position bias override (safe for rel_pos_arr=None) ───────────────────
+
+    def get_position_attn_bias(self, rel_pos_arr, seq_len, batch_size, device):
+        # When rel_pos_arr is None (no POS-aware dictionary available),
+        # skip that term rather than crashing.
+        tupe_bias = self.get_tupe_rel_pos_bias(seq_len, device) if self.use_tupe_rel_pos_bias else None
+
+        weight = self.pos_ln(self.pos.weight[:seq_len + 1, :])
+        pos_q = self.pos_q_linear(weight).view(seq_len + 1, self.seq_tr_nhead, -1).transpose(0, 1) * self.pos_scaling
+        pos_k = self.pos_k_linear(weight).view(seq_len + 1, self.seq_tr_nhead, -1).transpose(0, 1)
+        abs_pos_bias = torch.bmm(pos_q, pos_k.transpose(1, 2))
+
+        cls_2_other = abs_pos_bias[:, 0, 0]
+        other_2_cls = abs_pos_bias[:, 1, 1]
+        abs_pos_bias = abs_pos_bias[:, 1:, 1:]
+        abs_pos_bias[:, :, 0] = other_2_cls.view(-1, 1)
+        abs_pos_bias[:, 0, :] = cls_2_other.view(-1, 1)
+
+        if tupe_bias is not None:
+            abs_pos_bias = abs_pos_bias + tupe_bias
+
+        abs_pos_bias = abs_pos_bias.unsqueeze(0).expand(batch_size, -1, -1, -1).reshape(-1, seq_len, seq_len)
+
+        if self.use_pos_aware_rel_pos_bias and rel_pos_arr is not None:
+            abs_pos_bias = abs_pos_bias + self.get_pos_aware_rel_pos_bias(rel_pos_arr, seq_len)
+
+        return abs_pos_bias
+
+    # ── Table structure bias ──────────────────────────────────────────────────
+
+    def get_table_structure_bias(
+        self,
+        row_ids_padded: torch.Tensor,   # (N, S) — padded row ids
+        col_ids_padded: torch.Tensor,   # (N, S)
+        seq_len: int,
+        batch_size: int,
+    ) -> torch.Tensor:
+        """
+        Returns (N*H, S, S) tensor to be added to the position attn_bias.
+        row_ids_padded == 1 is the header row.
+        """
+        H = self.seq_tr_nhead
+        # (N, S, S) boolean: do positions i and j share the same row / col?
+        row_match = (row_ids_padded.unsqueeze(2) == row_ids_padded.unsqueeze(1)).float()
+        col_match = (col_ids_padded.unsqueeze(2) == col_ids_padded.unsqueeze(1)).float()
+        # Header: either token i or token j is in the header row (row_id == 1)
+        is_header = (row_ids_padded == 1)   # (N, S)
+        header_match = (is_header.unsqueeze(2) | is_header.unsqueeze(1)).float()  # (N, S, S)
+
+        # Broadcast head-specific scalars: (H,) → (H, 1, 1)
+        rb = self.row_attn_bias.view(H, 1, 1)
+        cb = self.col_attn_bias.view(H, 1, 1)
+        hb = self.header_attn_bias.view(H, 1, 1)
+
+        # (N, 1, S, S) * (H, 1, 1) broadcast → (N, H, S, S)
+        bias = (
+            rb * row_match.unsqueeze(1) +
+            cb * col_match.unsqueeze(1) +
+            hb * header_match.unsqueeze(1)
+        )
+        return bias.reshape(batch_size * H, seq_len, seq_len).contiguous()
+
+    # ── Forward ───────────────────────────────────────────────────────────────
+
+    @custom_fwd
+    def forward(self, args, rel_pos_arr,
+                tokens_lengths, input_sequence_lengths,
+                pos_tags, stems, afsets, affixes,
+                row_ids=None, col_ids=None, cell_types=None):
+        """
+        All parameters identical to KinyaBERT_MorphoEncoder.forward() plus:
+          row_ids, col_ids, cell_types — flat tensors of length L (total tokens).
+        """
+        device = stems.device
+
+        # ── Tier 1: Morpho transformer (unchanged from KinyaBERT) ─────────────
+        x_embed = None
+        if self.num_pos_m_embeddings > 0:
+            xm_pos1 = self.m1_pos_embedding(pos_tags)
+            x_embed = torch.unsqueeze(xm_pos1, 0)
+        if self.num_pos_m_embeddings > 1:
+            xm_pos2 = self.m2_pos_embedding(pos_tags)
+            x_embed = torch.cat((x_embed, torch.unsqueeze(xm_pos2, 0)), 0)
+        if self.num_pos_m_embeddings > 2:
+            xm_pos3 = self.m3_pos_embedding(pos_tags)
+            x_embed = torch.cat((x_embed, torch.unsqueeze(xm_pos3, 0)), 0)
+        if self.num_stem_m_embeddings > 0:
+            xm_stem = self.m_stem_embedding(stems)
+            xm_stem = torch.unsqueeze(xm_stem, 0)
+            x_embed = torch.cat((x_embed, xm_stem), 0) if x_embed is not None else xm_stem
+        if args.use_afsets and afsets is not None:
+            xm_afset = self.m_afset_embedding(afsets)
+            xm_afset = torch.unsqueeze(xm_afset, 0)
+            x_embed = torch.cat((x_embed, xm_afset), 0) if x_embed is not None else xm_afset
+
+        m_masks_padded = None
+        has_morphemes = False
+        morpho_input = None
+
+        if args.use_morpho_encoder:
+            afx = [a for a in affixes.split(tokens_lengths) if a.numel() > 0]
+            afx_padded = pad_sequence(afx, batch_first=False) if afx else torch.empty(0, device=device)
+            if afx_padded.nelement() > 0:
+                has_morphemes = True
+                xm_affix = self.m_affix_embedding(afx_padded)
+                x_embed = torch.cat((x_embed, xm_affix), 0) if x_embed is not None else xm_affix
+                m_masks = [torch.zeros(x + self.tot_morpho_idx, dtype=torch.bool, device=device)
+                           for x in tokens_lengths]
+                m_masks_padded = pad_sequence(m_masks, batch_first=True, padding_value=1)
+
+        if x_embed is not None and args.use_morpho_encoder:
+            morpho_out = self.morpho_transformer_encoder(x_embed, src_key_padding_mask=m_masks_padded)
+            if self.tot_morpho_idx > 0 and self.use_affix_bow_m_embedding:
+                heads = morpho_out[:self.tot_morpho_idx]
+                bow = torch.sum(morpho_out[self.tot_morpho_idx:], 0, keepdim=True) if has_morphemes \
+                    else torch.zeros((1, stems.size(0), self.morpho_dim), device=device)
+                morpho_input = torch.cat((heads, bow), 0)
+            elif self.tot_morpho_idx > 0:
+                morpho_input = morpho_out[:self.tot_morpho_idx]
+            elif self.use_affix_bow_m_embedding:
+                morpho_input = torch.sum(morpho_out[self.tot_morpho_idx:], 0, keepdim=True) if has_morphemes \
+                    else torch.zeros((1, stems.size(0), self.morpho_dim), device=device)
+        elif self.use_affix_bow_m_embedding:
+            morpho_input = torch.zeros((1, stems.size(0), self.morpho_dim), device=device)
+
+        # ── Build input to sequence transformer ──────────────────────────────
+        input_sequences = self.s_stem_embedding(stems)   # (L, stem_dim)
+
+        if morpho_input is not None:
+            morpho_input = morpho_input.permute(1, 0, 2).contiguous()
+            L = morpho_input.size(0)
+            morpho_input = morpho_input.view(L, -1)
+            input_sequences = torch.cat((morpho_input, input_sequences), 1)   # (L, seq_tr_d_model)
+
+        lists = input_sequences.split(input_sequence_lengths, 0)
+        tr_padded = pad_sequence(lists, batch_first=False)   # (S, N, d)
+
+        seq_len = tr_padded.size(0)
+        batch_size = tr_padded.size(1)
+
+        # ── Tabular additive embeddings ───────────────────────────────────────
+        row_ids_padded = None
+        col_ids_padded = None
+        if row_ids is not None:
+            row_lists = row_ids.split(input_sequence_lengths, 0)
+            row_ids_padded = pad_sequence(row_lists, batch_first=True, padding_value=0)  # (N, S)
+            col_lists = col_ids.split(input_sequence_lengths, 0)
+            col_ids_padded = pad_sequence(col_lists, batch_first=True, padding_value=0)
+            ct_lists = cell_types.split(input_sequence_lengths, 0)
+            ct_padded = pad_sequence(ct_lists, batch_first=True, padding_value=0)
+
+            # (N, S, d) → (S, N, d)
+            row_embed = self.row_embedding(row_ids_padded).permute(1, 0, 2)
+            col_embed = self.col_embedding(col_ids_padded).permute(1, 0, 2)
+            ct_embed  = self.cell_type_embedding(ct_padded).permute(1, 0, 2)
+            tr_padded = tr_padded + row_embed + col_embed + ct_embed
+
+        # ── Position bias (TUPE + POS-aware) — unchanged ─────────────────────
+        abs_pos_bias = self.get_position_attn_bias(rel_pos_arr, seq_len, batch_size, device)
+
+        # ── Table structure bias ──────────────────────────────────────────────
+        if row_ids_padded is not None and not getattr(args, 'no_bias', False):
+            table_bias = self.get_table_structure_bias(
+                row_ids_padded, col_ids_padded, seq_len, batch_size
+            )
+            abs_pos_bias = abs_pos_bias + table_bias
+
+        # ── Tier 2: Sequence transformer ──────────────────────────────────────
+        masks = [torch.zeros(x, dtype=torch.bool, device=device) for x in input_sequence_lengths]
+        masks_padded = pad_sequence(masks, batch_first=True, padding_value=1)
+
+        transformer_output = self.seq_transformer_encoder(
+            tr_padded,
+            attn_bias=abs_pos_bias,
+            src_key_padding_mask=masks_padded,
+        )   # (S, N, d)
+
+        return transformer_output
+
+    # ── CTP prediction ────────────────────────────────────────────────────────
+
+    def predict_ctp(self, hidden_states: torch.Tensor,
+                    ctp_positions: torch.Tensor) -> torch.Tensor:
+        """
+        hidden_states: (S, N, d)
+        ctp_positions: (num_ctp,) — flat indices into the (N*S) token space
+        Returns: (num_ctp, CTP_NUM_LABELS) logits
+        """
+        flat = hidden_states.permute(1, 0, 2).reshape(-1, self.seq_tr_d_model)
+        selected = torch.index_select(flat, 0, ctp_positions)
+        return self.ctp_head(selected)
+
+
+# ── TabuLM (full model with all prediction heads) ─────────────────────────────
+
+class TabuLM(nn.Module):
+    """Full TabuLM pre-training model."""
+
+    def __init__(self, args,
+                 num_stems, num_afsets, num_pos_tags, num_affixes,
+                 num_rel_pos_dict_size,
+                 num_pos_m_embeddings, num_stem_m_embeddings,
+                 use_affix_bow_m_embedding,
+                 use_pos_aware_rel_pos_bias, use_tupe_rel_pos_bias,
+                 max_seq_len=512,
+                 morpho_dim=128, stem_dim=256,
+                 morpho_tr_nhead=4, morpho_tr_nlayers=4,
+                 morpho_tr_dim_feedforward=512, morpho_tr_dropout=0.1,
+                 morpho_tr_activation='gelu',
+                 seq_tr_nhead=12, seq_tr_nlayers=12,
+                 seq_tr_dim_feedforward=3072, seq_tr_dropout=0.1,
+                 seq_tr_activation='gelu',
+                 layernorm_epsilon=1e-6,
+                 tupe_rel_pos_bins=32, tupe_max_rel_pos=128,
+                 max_rows=MAX_ROWS, max_cols=MAX_COLS):
+        super().__init__()
+
+        self.encoder = TabuLM_Encoder(
+            args,
+            num_stems, num_afsets, num_pos_tags, num_affixes,
+            num_rel_pos_dict_size,
+            num_pos_m_embeddings, num_stem_m_embeddings,
+            use_affix_bow_m_embedding,
+            use_pos_aware_rel_pos_bias, use_tupe_rel_pos_bias,
+            max_seq_len=max_seq_len,
+            morpho_dim=morpho_dim, stem_dim=stem_dim,
+            morpho_tr_nhead=morpho_tr_nhead, morpho_tr_nlayers=morpho_tr_nlayers,
+            morpho_tr_dim_feedforward=morpho_tr_dim_feedforward,
+            morpho_tr_dropout=morpho_tr_dropout, morpho_tr_activation=morpho_tr_activation,
+            seq_tr_nhead=seq_tr_nhead, seq_tr_nlayers=seq_tr_nlayers,
+            seq_tr_dim_feedforward=seq_tr_dim_feedforward,
+            seq_tr_dropout=seq_tr_dropout, seq_tr_activation=seq_tr_activation,
+            tupe_rel_pos_bins=tupe_rel_pos_bins, tupe_max_rel_pos=tupe_max_rel_pos,
+            max_rows=max_rows, max_cols=max_cols,
+        )
+
+        d = self.encoder.seq_tr_d_model
+
+        from morpho_model import MorphoHeadPredictor
+        self.head_predictor = MorphoHeadPredictor(
+            args,
+            stem_embedding_weights=self.encoder.s_stem_embedding.weight,
+            afset_embedding_weights=(
+                self.encoder.m_afset_embedding.weight
+                if hasattr(self.encoder, 'm_afset_embedding') else
+                torch.zeros(1, morpho_dim)
+            ),
+            affix_embedding_weights=(
+                self.encoder.m_affix_embedding.weight
+                if hasattr(self.encoder, 'm_affix_embedding') else
+                torch.zeros(1, morpho_dim)
+            ),
+            tr_d_model=d,
+            tr_dropout=seq_tr_dropout,
+            layernorm_epsilon=layernorm_epsilon,
+        )
+
+    def forward(self, args, rel_pos_arr,
+                tokens_lengths, input_sequence_lengths,
+                pos_tags, stems, afsets, affixes,
+                row_ids, col_ids, cell_types,
+                predicted_tokens_idx, predicted_tokens_affixes_idx,
+                predicted_stems, predicted_afsets, predicted_affixes_prob,
+                mcr_tokens_idx=None, mcr_stems=None,
+                ctp_positions=None, ctp_labels=None):
+
+        hidden = self.encoder(
+            args, rel_pos_arr, tokens_lengths, input_sequence_lengths,
+            pos_tags, stems, afsets, affixes,
+            row_ids=row_ids, col_ids=col_ids, cell_types=cell_types,
+        )
+
+        # ── Standard KinyaBERT losses (stem + afset + affix) ──────────────────
+        n_stems = self.head_predictor.stem_decoder.weight.size(0)
+        predicted_stems = predicted_stems.clamp(0, n_stems - 1)
+        loss, stem_loss, afset_loss, affix_loss = self.head_predictor(
+            args, hidden,
+            predicted_tokens_idx, predicted_tokens_affixes_idx,
+            predicted_stems, predicted_afsets, predicted_affixes_prob,
+        )
+
+        # ── MCR loss — reuse stem decoder on cell-masked tokens ───────────────
+        mcr_loss = torch.tensor(0.0, device=loss.device)
+        if mcr_tokens_idx is not None and mcr_stems is not None and mcr_tokens_idx.numel() > 0:
+            flat = hidden.permute(1, 0, 2).reshape(-1, hidden.size(-1))
+            mcr_hidden = torch.index_select(flat, 0, mcr_tokens_idx)
+            mcr_state = self.head_predictor.stem_transform(mcr_hidden)
+            mcr_scores = self.head_predictor.stem_decoder(mcr_state) + self.head_predictor.stem_decoder_bias
+            n_stem_cls = mcr_scores.size(1)
+            mcr_stems_clamped = mcr_stems.clamp(0, n_stem_cls - 1)
+            mcr_loss = F.nll_loss(F.log_softmax(mcr_scores, dim=1), mcr_stems_clamped)
+            loss = loss + mcr_loss
+
+        # ── CTP loss ──────────────────────────────────────────────────────────
+        ctp_loss = torch.tensor(0.0, device=loss.device)
+        if ctp_positions is not None and ctp_labels is not None and ctp_positions.numel() > 0:
+            ctp_logits = self.encoder.predict_ctp(hidden, ctp_positions)
+            n_ctp_cls = ctp_logits.size(1)
+            ctp_labels_clamped = ctp_labels.clamp(0, n_ctp_cls - 1)
+            ctp_loss = F.cross_entropy(ctp_logits, ctp_labels_clamped)
+            loss = loss + ctp_loss
+
+        return loss, stem_loss, afset_loss, affix_loss, mcr_loss, ctp_loss
+
+    def state_dict(self, *args, **kwargs):
+        return super().state_dict(*args, **kwargs)
+
+    def load_state_dict(self, *args, **kwargs):
+        return super().load_state_dict(*args, **kwargs)
+
+
+# ── Factory ───────────────────────────────────────────────────────────────────
+
+def tabulm_base(kb_vocab, affix_set_vocab, morpho_rel_pos_dict,
+                device, args, saved_model_file=None):
+    """
+    Construct TabuLM-base model.  Mirrors kinyabert_base() from morpho_model.py.
+    All default hyperparameters match KinyaBERT-base so a KinyaBERT checkpoint
+    can be used to warm-start the encoder (tabular components start from scratch).
+    """
+    num_stems = len(kb_vocab.reduced_stem_vocab) + 1
+    num_afsets = (len(affix_set_vocab.affix_set_vocab_idx) + 1) if affix_set_vocab is not None else 0
+    num_pos_tags = len(kb_vocab.pos_tag_vocab) + 1
+    num_affixes = len(kb_vocab.affix_vocab) + 1
+    num_rel_pos_dict_size = (len(morpho_rel_pos_dict) + 1) if morpho_rel_pos_dict is not None else 0
+
+    model = TabuLM(
+        args,
+        num_stems=num_stems,
+        num_afsets=num_afsets,
+        num_pos_tags=num_pos_tags,
+        num_affixes=num_affixes,
+        num_rel_pos_dict_size=num_rel_pos_dict_size,
+        num_pos_m_embeddings=args.num_pos_m_embeddings,
+        num_stem_m_embeddings=args.num_stem_m_embeddings,
+        use_affix_bow_m_embedding=args.use_affix_bow_m_embedding,
+        use_pos_aware_rel_pos_bias=args.use_pos_aware_rel_pos_bias,
+        use_tupe_rel_pos_bias=args.use_tupe_rel_pos_bias,
+        max_seq_len=args.max_seq_len,
+        morpho_dim=args.morpho_dim,
+        stem_dim=args.stem_dim,
+        morpho_tr_nhead=args.morpho_tr_nhead,
+        morpho_tr_nlayers=args.morpho_tr_nlayers,
+        morpho_tr_dim_feedforward=args.morpho_tr_dim_feedforward,
+        morpho_tr_dropout=args.morpho_tr_dropout,
+        seq_tr_nhead=args.seq_tr_nhead,
+        seq_tr_nlayers=args.seq_tr_nlayers,
+        seq_tr_dim_feedforward=args.seq_tr_dim_feedforward,
+        seq_tr_dropout=args.seq_tr_dropout,
+        layernorm_epsilon=args.layernorm_epsilon,
+    ).to(device)
+
+    if saved_model_file is not None:
+        state = torch.load(saved_model_file, map_location=device)
+        # Warm-start from KinyaBERT checkpoint (ignore tabular-specific keys)
+        own_state = model.state_dict()
+        pretrained = state.get('model_state_dict', state)
+        matched, skipped = 0, 0
+        for k, v in pretrained.items():
+            # Map KinyaBERT encoder keys to TabuLM encoder keys
+            new_k = k.replace('module.encoder.', 'encoder.')
+            if new_k not in own_state:
+                new_k = k
+            if new_k in own_state and own_state[new_k].shape == v.shape:
+                own_state[new_k].copy_(v)
+                matched += 1
+            else:
+                skipped += 1
+        print(f'[tabulm_base] warm-started {matched} layers, skipped {skipped}')
+        model.load_state_dict(own_state)
+
+    return model

code/train_tabulm.py

@@ -0,0 +1,379 @@
+# TabuLM — pre-training script
+# Extends train_exploratory_distributed_model.py for tabular data.
+# Logs: STEM | AFSET | AFFIX | MCR | CTP losses separately.
+
+from __future__ import print_function, division
+
+import gc
+import math
+import os
+import random
+from datetime import datetime
+from shutil import copyfile
+
+import numpy as np
+import progressbar
+import psutil
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.data import DataLoader
+
+
+def time_now():
+    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+
+def date_now():
+    return datetime.now().strftime("%Y-%m-%d")
+
+def set_random_seeds(seed=0):
+    torch.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+    np.random.seed(seed)
+    random.seed(seed)
+
+
+def train_loop(args, rank, scaler, device, data_loader,
+               model, optimizer, lr_scheduler, save_file_path,
+               accumulation_steps, loop, num_loops, bar,
+               total_steps, total_loss,
+               stem_loss_acc, afset_loss_acc, affix_loss_acc,
+               mcr_loss_acc, ctp_loss_acc,
+               save_every=50):
+
+    from tabular_data_loaders import tabulm_model_forward
+
+    for batch_idx, data_item in enumerate(data_loader):
+        if scaler is not None:
+            with torch.cuda.amp.autocast():
+                loss, sl, al, fxl, ml, cl = tabulm_model_forward(
+                    args, data_item, model, device,
+                    model.module.encoder.tot_num_affixes if hasattr(model, 'module')
+                    else model.encoder.tot_num_affixes,
+                )
+                loss = loss / accumulation_steps
+            scaler.scale(loss).backward()
+        else:
+            loss, sl, al, fxl, ml, cl = tabulm_model_forward(
+                args, data_item, model, device,
+                model.module.encoder.tot_num_affixes if hasattr(model, 'module')
+                else model.encoder.tot_num_affixes,
+            )
+            loss = loss / accumulation_steps
+            loss.backward()
+
+        total_loss      += loss.item()
+        stem_loss_acc   += sl.item()  / accumulation_steps
+        afset_loss_acc  += al.item()  / accumulation_steps
+        affix_loss_acc  += fxl.item() / accumulation_steps
+        mcr_loss_acc    += ml.item()  / accumulation_steps
+        ctp_loss_acc    += cl.item()  / accumulation_steps
+        total_steps     += 1
+
+        if (total_steps % accumulation_steps) == 0:
+            if scaler is not None:
+                scaler.step(optimizer)
+                scaler.update()
+            else:
+                optimizer.step()
+            optimizer.zero_grad()
+            lr_scheduler.step()
+
+            if rank == 0:
+                print(
+                    time_now(),
+                    f'Loop:{loop}/{num_loops}',
+                    f'Batch:{batch_idx+1}/{len(data_loader)}',
+                    f'TOTAL:{total_loss:.4f}',
+                    f'STEM:{stem_loss_acc:.4f}',
+                    f'AFSET:{afset_loss_acc:.4f}',
+                    f'AFFIX:{affix_loss_acc:.4f}',
+                    f'MCR:{mcr_loss_acc:.4f}',
+                    f'CTP:{ctp_loss_acc:.4f}',
+                    f'LR:{lr_scheduler.get_lr():.8f}',
+                    f'iter:{lr_scheduler.num_iters}',
+                )
+                bar.update(lr_scheduler.num_iters)
+
+            total_loss = stem_loss_acc = afset_loss_acc = 0.0
+            affix_loss_acc = mcr_loss_acc = ctp_loss_acc = 0.0
+
+    if rank == 0 and (((loop + 1) % save_every) == 0 or loop == num_loops - 1):
+        if os.path.exists(save_file_path):
+            copyfile(save_file_path, save_file_path + '_prev_checkpoint.pt')
+
+        _model = model.module if hasattr(model, 'module') else model
+        _model.eval()
+        torch.save({
+            'iter': total_steps,
+            'model_state_dict': model.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            'lr_scheduler_state_dict': lr_scheduler.state_dict(),
+            'loop': loop,
+            'num_loops': num_loops,
+        }, save_file_path)
+        _model.train()
+
+    return (total_steps, total_loss,
+            stem_loss_acc, afset_loss_acc, affix_loss_acc,
+            mcr_loss_acc, ctp_loss_acc)
+
+
+def train_fn(rank, args):
+    import youtokentome as yttm
+    from morpho_learning_rates import AnnealingLR
+    from morpho_data_loaders import KBVocab, AffixSetVocab
+    from tabular_data_loaders import TabularKBCorpusDataset, tabular_collate_wrapper
+    from tabulm_model import tabulm_base
+
+    USE_GPU = args.gpus > 0 and torch.cuda.is_available()
+
+    device = torch.device('cuda' if USE_GPU else 'cpu')
+
+    if USE_GPU:
+        dist.init_process_group('nccl', init_method='env://',
+                                world_size=args.world_size, rank=rank)
+        torch.cuda.set_device(rank)
+        scaler = torch.cuda.amp.GradScaler()
+    else:
+        dist.init_process_group('gloo', init_method='env://',
+                                world_size=args.world_size, rank=rank)
+        scaler = None
+
+    home = args.home_path
+
+    bpe = yttm.BPE(model=home + 'data/BPE-30k.mdl')
+
+    kb_vocab = KBVocab()
+    kb_vocab.load_state_dict(torch.load(home + 'data/kb_vocab_state_dict_2021-02-07.pt'))
+
+    affix_set_vocab = None
+    if args.use_afsets:
+        affix_set_vocab = AffixSetVocab(
+            reduced_affix_dict_file=home + f'data/reduced_affix_dict_{args.afset_dict_size}.csv',
+            reduced_affix_dict_map_file=home + f'data/reduced_affix_dict_map_{args.afset_dict_size}.csv',
+        )
+
+    morpho_rel_pos_dict = None
+    morpho_rel_pos_dmax = 5
+    if args.use_pos_aware_rel_pos_bias:
+        rel_pos_file = home + 'data/morpho_rel_pos_dict_2021-03-24.pt'
+        if os.path.exists(rel_pos_file):
+            saved = torch.load(rel_pos_file)
+            morpho_rel_pos_dict = saved['morpho_rel_pos_dict']
+            morpho_rel_pos_dmax = saved['morpho_rel_pos_dmax']
+        else:
+            print(f'[WARN] morpho_rel_pos_dict not found, disabling pos_aware_rel_pos_bias')
+            args.use_pos_aware_rel_pos_bias = False
+            args.use_pos_aware_rel = False
+
+    num_iters   = args.num_iters
+    warmup_iter = args.warmup_iter
+    peak_lr     = args.peak_lr
+    wd          = args.wd
+
+    if rank == 0:
+        print(time_now(), 'Building TabuLM model ...')
+
+    model = tabulm_base(kb_vocab, affix_set_vocab, morpho_rel_pos_dict,
+                        device, args, saved_model_file=args.exploratory_model_load)
+
+    if USE_GPU:
+        model = DDP(model, device_ids=[rank], find_unused_parameters=True)
+        try:
+            import apex
+            optimizer = apex.optimizers.FusedLAMB(
+                model.parameters(), lr=peak_lr, betas=(0.9, 0.98),
+                eps=1e-06, weight_decay=wd,
+            )
+        except ImportError:
+            from lamb import Lamb
+            optimizer = Lamb(model.parameters(), lr=peak_lr, betas=(0.9, 0.98),
+                             eps=1e-06, weight_decay=wd)
+    else:
+        from lamb import Lamb
+        model = DDP(model, device_ids=[])
+        optimizer = Lamb(model.parameters(), lr=peak_lr, betas=(0.9, 0.98),
+                         eps=1e-06, weight_decay=wd)
+
+    lr_scheduler = AnnealingLR(optimizer,
+                               start_lr=peak_lr,
+                               warmup_iter=warmup_iter,
+                               num_iters=num_iters,
+                               decay_style='linear',
+                               last_iter=0)
+
+    # ── Resume from checkpoint if provided ────────────────────────────────────
+    resume_file = getattr(args, 'resume_checkpoint', None)
+    curr_loops = 0
+    total_steps = 0
+    total_loss  = stem_loss_acc = afset_loss_acc = 0.0
+    affix_loss_acc = mcr_loss_acc = ctp_loss_acc = 0.0
+
+    if resume_file and os.path.exists(resume_file):
+        if rank == 0:
+            print(f'[RESUME] Loading checkpoint from {resume_file}')
+        ckpt = torch.load(resume_file, map_location=device)
+        # Strip DDP 'module.' prefix if present
+        state = ckpt['model_state_dict']
+        if all(k.startswith('module.') for k in state):
+            state = {k[len('module.'):]: v for k, v in state.items()}
+        _model = model.module if hasattr(model, 'module') else model
+        _model.load_state_dict(state, strict=False)
+        optimizer.load_state_dict(ckpt['optimizer_state_dict'])
+        lr_scheduler.load_state_dict(ckpt['lr_scheduler_state_dict'])
+        curr_loops  = ckpt.get('loop', 0) + 1
+        total_steps = ckpt.get('iter', 0)
+        if rank == 0:
+            print(f'[RESUME] Resuming from loop {curr_loops}, iter {total_steps}')
+
+    csv_dir = args.tabulm_csv_dir if hasattr(args, 'tabulm_csv_dir') and args.tabulm_csv_dir \
+        else home + 'data/tables/'
+
+    num_train_loops = math.ceil(
+        num_iters * args.accumulation_steps / args.number_of_load_batches
+    )
+
+    save_path = (
+        home + f'data/tabulm_model_{date_now()}'
+        f'_pos@{args.num_pos_m_embeddings}'
+        f'_stem@{args.num_stem_m_embeddings}'
+        f'_afsets@{args.use_afsets}@{args.afset_dict_size}'
+        f'{getattr(args, "ablation_tag", "")}.pt'
+    )
+
+    total_params     = sum(p.numel() for p in model.parameters())
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+    if rank == 0:
+        print('─' * 60)
+        print(f'Total params: {total_params:,}  Trainable: {trainable_params:,}')
+        print(f'Saving to: {save_path}')
+        print(f'CSV tables from: {csv_dir}')
+        print(f'num_iters={num_iters}  warmup={warmup_iter}  loops={num_train_loops}')
+        print(f'batch_size={args.batch_size}  accum={args.accumulation_steps}')
+        print(f'peak_lr={peak_lr}  wd={wd}')
+        print('─' * 60)
+
+    model.train()
+    model.zero_grad()
+
+    with progressbar.ProgressBar(
+        initial_value=lr_scheduler.num_iters,
+        max_value=lr_scheduler.end_iter,
+        redirect_stdout=True,
+    ) as bar:
+        if rank == 0:
+            bar.update(lr_scheduler.num_iters)
+
+        for loop in range(curr_loops, num_train_loops):
+            if rank == 0:
+                print(time_now(), 'Loading tabular dataset ...')
+
+            dataset = TabularKBCorpusDataset(
+                args, kb_vocab, affix_set_vocab, bpe,
+                csv_dir=csv_dir,
+                max_batch_items=args.number_of_load_batches * args.batch_size,
+                max_seq_len=512,
+                rank=rank,
+            )
+
+            data_loader = DataLoader(
+                dataset,
+                batch_size=args.batch_size,
+                collate_fn=tabular_collate_wrapper,
+                shuffle=True,
+            )
+
+            if rank == 0:
+                print(time_now(), 'Memory:', psutil.virtual_memory())
+
+            (total_steps, total_loss,
+             stem_loss_acc, afset_loss_acc, affix_loss_acc,
+             mcr_loss_acc, ctp_loss_acc) = train_loop(
+                args, rank, scaler, device, data_loader,
+                model, optimizer, lr_scheduler, save_path,
+                args.accumulation_steps, loop, num_train_loops, bar,
+                total_steps, total_loss,
+                stem_loss_acc, afset_loss_acc, affix_loss_acc,
+                mcr_loss_acc, ctp_loss_acc,
+                save_every=getattr(args, 'save_every', 50),
+            )
+
+            if rank == 0:
+                print(time_now(), f'{loop+1}/{num_train_loops} loops complete')
+
+            del data_loader, dataset
+            gc.collect()
+
+
+def main():
+    import argparse
+    from morpho_common import setup_common_args
+
+    # Pull out --resume-checkpoint and ablation flags before setup_common_args sees sys.argv
+    import sys
+    resume_checkpoint = None
+    no_mcr = False
+    no_ctp = False
+    no_tabular_emb = False
+    no_bias = False
+    ablation_tag = ''
+    filtered = []
+    i = 0
+    while i < len(sys.argv[1:]):
+        arg = sys.argv[1:][i]
+        if arg == '--resume-checkpoint':
+            resume_checkpoint = sys.argv[1:][i + 1]
+            i += 2
+        elif arg.startswith('--resume-checkpoint='):
+            resume_checkpoint = arg.split('=', 1)[1]
+            i += 1
+        elif arg == '--no-mcr':
+            no_mcr = True
+            ablation_tag += '_noMCR'
+            i += 1
+        elif arg == '--no-ctp':
+            no_ctp = True
+            ablation_tag += '_noCTP'
+            i += 1
+        elif arg == '--no-tabular-emb':
+            no_tabular_emb = True
+            ablation_tag += '_noTabEmb'
+            i += 1
+        elif arg == '--no-bias':
+            no_bias = True
+            ablation_tag += '_noBias'
+            i += 1
+        else:
+            filtered.append(arg)
+            i += 1
+    sys.argv = [sys.argv[0]] + filtered
+
+    args = setup_common_args()
+    args.resume_checkpoint = resume_checkpoint
+    args.no_mcr = no_mcr
+    args.no_ctp = no_ctp
+    args.no_tabular_emb = no_tabular_emb
+    args.no_bias = no_bias
+    args.ablation_tag = ablation_tag
+
+    # Extra args not in morpho_common.setup_common_args
+    if not hasattr(args, 'tabulm_csv_dir'):
+        args.tabulm_csv_dir = os.environ.get('TABULM_CSV_DIR', None)
+    if not hasattr(args, 'resume_checkpoint'):
+        args.resume_checkpoint = None
+
+    os.environ['MASTER_ADDR'] = 'localhost'
+    os.environ['MASTER_PORT'] = os.environ.get('MASTER_PORT', '29602')
+
+    if args.gpus == 0:
+        args.world_size = 1
+
+    mp.spawn(train_fn, nprocs=args.world_size, args=(args,))
+
+
+if __name__ == '__main__':
+    main()