code/finetune_tableqa.py

"""
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()