train.py

"""
TELEN: Temporal Evolving Legal Embedding Network — Training Script.

Stages:
  1. Contrastive pretraining (5 epochs) — train projection head
  2. Meta-training (50 epochs) — train HyperNetwork + State Encoder

Usage:
    python train.py
"""
import sys, os, math, random
from pathlib import Path
from collections import defaultdict
import numpy as np
import pandas as pd
import torch, torch.nn as nn, torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm
from transformers import AutoTokenizer
from pyvi import ViTokenizer

sys.path.insert(0, ".")
from src.telern.config import TELENConfig, DATA_DIR
from src.telern.model import TELEN, create_model
from src.data import load_raw_data, extract_metadata, clean_data

SEED = 42
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ═══════════════════════════════════════════════════════════
# Data
# ═══════════════════════════════════════════════════════════
def prepare_data(config):
    df = load_raw_data(str(DATA_DIR / "train-00000-of-00001.parquet"))
    df = extract_metadata(df); df = clean_data(df, min_text_len=10)
    articles_by_law = defaultdict(list)
    laws_by_year = defaultdict(list)
    for _, row in df.iterrows():
        articles_by_law[row["law_id"]].append({
            "id": row["id"], "title": row["title"], "text": row["text"],
            "law_type": row["law_type"], "year": row["year"],
        })
    for law_id in articles_by_law:
        laws_by_year[articles_by_law[law_id][0]["year"]].append(law_id)
    all_years = sorted(laws_by_year.keys())
    train_years = [y for y in all_years if y <= config.meta.train_split_year]
    val_years = [y for y in all_years if config.meta.train_split_year < y <= config.meta.val_split_year]
    test_years = [y for y in all_years if y > config.meta.val_split_year]
    return articles_by_law, laws_by_year, train_years, val_years, test_years, df

# ═══════════════════════════════════════════════════════════
# Contrastive Dataset
# ═══════════════════════════════════════════════════════════
class ContrastiveDataset(Dataset):
    def __init__(self, df, tokenizer, max_len=480):
        self.df = df.reset_index(drop=True)
        self.tokenizer = tokenizer
        self.max_len = max_len
        self.law_groups = self.df.groupby("law_id")
        self.law_ids = list(self.law_groups.groups.keys())

    def __len__(self): return len(self.df)

    def __getitem__(self, idx):
        row = self.df.iloc[idx]; law_id = row["law_id"]
        wseg = lambda t: ViTokenizer.tokenize(t.replace("_", " "))
        anchor = wseg(f"{row['title']}: {row['text'][:400]}")
        group_idx = self.law_groups.groups[law_id]
        others = [i for i in group_idx if i != idx]
        pos_row = self.df.iloc[random.choice(others)] if others else row
        positive = wseg(f"{pos_row['title']}: {pos_row['text'][:400]}")
        neg_law = random.choice([l for l in self.law_ids if l != law_id])
        neg_row = self.df.iloc[random.choice(list(self.law_groups.groups[neg_law]))]
        negative = wseg(f"{neg_row['title']}: {neg_row['text'][:400]}")

        def tok(t): return self.tokenizer(t, truncation=True, max_length=self.max_len, padding="max_length", return_tensors="pt")
        return {f"{k}_{s}": tok(t)[k].squeeze(0)
                for t, s in [(anchor,"a"),(positive,"p"),(negative,"n")]
                for k in ["input_ids","attention_mask"]}

# ═══════════════════════════════════════════════════════════
# Stage 1: Contrastive Pretraining
# ═══════════════════════════════════════════════════════════
def contrastive_pretrain(model, df, config, epochs=5, batch_size=24, lr=3e-5):
    tokenizer = model.encoder.tokenizer
    dataset = ContrastiveDataset(df, tokenizer, config.max_seq_length)
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

    trainable = list(model.base_projection.parameters()) + [model.attn_query]
    opt = torch.optim.AdamW(trainable, lr=lr, weight_decay=0.01)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs * len(loader))

    print(f"  Contrastive pretraining: {epochs} epochs, {len(loader)} batches")
    model.train(); model.encoder.model.eval()

    for epoch in range(epochs):
        total = 0.0
        for batch in tqdm(loader, desc=f"  Epoch {epoch+1}/{epochs}"):
            a_ids=batch["input_ids_a"].to(device); a_mask=batch["attention_mask_a"].to(device)
            p_ids=batch["input_ids_p"].to(device); p_mask=batch["attention_mask_p"].to(device)
            n_ids=batch["input_ids_n"].to(device); n_mask=batch["attention_mask_n"].to(device)

            with torch.no_grad():
                ah=model._pool(model.encoder.model(input_ids=a_ids,attention_mask=a_mask).last_hidden_state,a_mask)
                ph=model._pool(model.encoder.model(input_ids=p_ids,attention_mask=p_mask).last_hidden_state,p_mask)
                nh=model._pool(model.encoder.model(input_ids=n_ids,attention_mask=n_mask).last_hidden_state,n_mask)

            ae=F.normalize(model.base_projection(ah),p=2,dim=1)
            pe=F.normalize(model.base_projection(ph),p=2,dim=1)
            ne=F.normalize(model.base_projection(nh),p=2,dim=1)

            trip=F.relu(0.3-(ae*pe).sum(1)+(ae*ne).sum(1)).mean()
            sim=ae@torch.cat([ae,pe,ne],dim=0).T/0.05
            infonce=F.cross_entropy(sim,torch.arange(len(a_ids),device=device)+len(a_ids))
            loss=trip+0.5*infonce

            opt.zero_grad(); loss.backward()
            torch.nn.utils.clip_grad_norm_(trainable,1.0); opt.step(); sched.step()
            total+=loss.item()
        print(f"    Epoch {epoch+1} avg loss: {total/len(loader):.4f}")
    print("  Contrastive pretraining complete!")
    return model

# ═══════════════════════════════════════════════════════════
# Episode building
# ═══════════════════════════════════════════════════════════
def build_episode(articles_by_law, laws_by_year, state_years, query_year, config):
    mc = config.meta
    q_laws = laws_by_year.get(query_year, [])
    if len(q_laws) < 5: return None
    sampled = random.sample(q_laws, min(mc.n_query // 4, len(q_laws)))
    queries, positives, q_types = [], [], set()
    for lid in sampled:
        arts = articles_by_law[lid]
        if len(arts) < 2: continue
        qi, pi = random.sample(range(len(arts)), 2)
        queries.append(arts[qi]); positives.append(arts[pi])
        q_types.add(arts[qi]["law_type"])
    if len(queries) < 4: return None

    hard_neg, rand_neg = [], []
    for lid in q_laws:
        if lid in sampled: continue
        for a in articles_by_law[lid]:
            if a["law_type"] in q_types: hard_neg.append(a)
            else: rand_neg.append(a)

    nh = min(mc.n_negatives // 2, len(hard_neg))
    nr = min(mc.n_negatives - nh, len(rand_neg))
    negatives = (random.sample(hard_neg, nh) if nh > 0 else []) + (random.sample(rand_neg, nr) if nr > 0 else [])
    if len(negatives) < 4: return None
    return {"queries": queries, "positives": positives, "negatives": negatives}

# ═══════════════════════════════════════════════════════════
# Stage 2: Meta-Training
# ═══════════════════════════════════════════════════════════
def compute_loss(model, q_texts, p_texts, n_texts, state_vec, temp=0.05):
    n_q, n_p = len(q_texts), len(p_texts)
    if n_q == 0 or n_p == 0:
        return torch.tensor(0.0, device=device, requires_grad=True)
    all_t = q_texts + p_texts + n_texts
    raw = model.encode_text(all_t)
    adapted = model.adapt_embedding(raw, state_vec)
    emb = adapted["mean"]
    qe = emb[:n_q]; pe = emb[n_q:n_q+n_p]; ne = emb[n_q+n_p:]

    if n_q == n_p:
        sim = torch.cat([(qe*pe).sum(1).unsqueeze(1)/temp, qe@ne.T/temp], dim=1)
        loss = F.cross_entropy(sim, torch.zeros(n_q, dtype=torch.long, device=device))
    else:
        loss = F.cross_entropy(qe @ torch.cat([pe, ne], dim=0).T / temp,
                               torch.arange(n_q, device=device).clamp(max=len(pe)-1))

    if model.config.hypernetwork.output_variance:
        lv = adapted.get("log_variance")
        if lv is not None: loss = loss + (lv.exp() - lv - 1).mean() * model.config.meta.kl_weight
    return loss

def validate(model, articles_by_law, laws_by_year, val_years, config):
    model.eval(); losses = []
    with torch.no_grad():
        for _ in range(30):
            qy = random.choice(val_years)
            if qy not in laws_by_year: continue
            sy = [y for y in sorted(laws_by_year.keys()) if y < qy]
            if len(sy) < 3: sy = [y for y in sorted(laws_by_year.keys()) if y <= qy]
            ep = build_episode(articles_by_law, laws_by_year, sy, qy, config)
            if ep is None: continue
            sv = model.get_state_vector()
            losses.append(compute_loss(model,
                [f"{q['title']}: {q['text'][:200]}" for q in ep["queries"]],
                [f"{p['title']}: {p['text'][:200]}" for p in ep["positives"]],
                [f"{n['title']}: {n['text'][:200]}" for n in ep["negatives"]],
                sv, config.meta.temperature).item())
    return sum(losses)/max(len(losses),1)

def meta_train(model, articles_by_law, laws_by_year, train_years, val_years, config, epochs=50, patience=15):
    trainable = (list(model.hypernetwork.parameters()) + list(model.state_encoder.parameters()) +
                 list(model.base_projection.parameters()) + [model.attn_query])
    opt = torch.optim.AdamW(trainable, lr=config.meta.meta_lr, weight_decay=1e-4)
    sched = torch.optim.lr_scheduler.ReduceLROnPlateau(opt, mode='min', factor=0.5, patience=5, min_lr=1e-6)

    os.makedirs(config.output_dir, exist_ok=True)
    best_val, patience_ctr = float("inf"), 0

    for epoch in range(epochs):
        model.train(); total_loss = 0.0
        steps = config.meta.meta_batch_size * 100
        progress = tqdm(range(steps), desc=f"Meta Epoch {epoch+1}/{epochs}")
        for _ in progress:
            if len(train_years) < 3: break
            si = random.randint(2, len(train_years)-1)
            sy, qy = train_years[:si], train_years[si]
            if qy not in laws_by_year: continue
            ep = build_episode(articles_by_law, laws_by_year, sy, qy, config)
            if ep is None: continue
            sv = model.get_state_vector()
            loss = compute_loss(model,
                [f"{q['title']}: {q['text'][:200]}" for q in ep["queries"]],
                [f"{p['title']}: {p['text'][:200]}" for p in ep["positives"]],
                [f"{n['title']}: {n['text'][:200]}" for n in ep["negatives"]],
                sv, config.meta.temperature)
            opt.zero_grad(); loss.backward()
            torch.nn.utils.clip_grad_norm_(trainable, 1.0); opt.step()
            total_loss += loss.item()
            progress.set_postfix({"loss": f"{loss.item():.4f}"})

        avg = total_loss / max(steps, 1)
        print(f"  Epoch {epoch+1} avg loss: {avg:.4f}")

        vl = validate(model, articles_by_law, laws_by_year, val_years, config)
        print(f"  Val loss: {vl:.4f}")
        sched.step(vl)

        if vl < best_val:
            best_val, patience_ctr = vl, 0
            torch.save({
                "hypernetwork": model.hypernetwork.state_dict(),
                "state_encoder": model.state_encoder.state_dict(),
                "base_projection": model.base_projection.state_dict(),
                "attn_query": model.attn_query,
                "epoch": epoch, "val_loss": vl,
            }, Path(config.output_dir) / "telen_best.pt")
            print(f"  Saved (val_loss={vl:.4f})")
        else:
            patience_ctr += 1
            if patience_ctr >= patience:
                print(f"  Early stopping at epoch {epoch+1}"); break
    print("Meta-training complete!")
    return model

# ═══════════════════════════════════════════════════════════
# Main
# ═══════════════════════════════════════════════════════════
def main():
    config = TELENConfig()
    random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
    print(f"Device: {device}")

    # Data
    print("\nLoading data...")
    articles_by_law, laws_by_year, train_years, val_years, test_years, df = prepare_data(config)
    print(f"  Train: {train_years[0]}-{train_years[-1]} ({len(train_years)}y)")
    print(f"  Val:   {val_years[0]}-{val_years[-1]} ({len(val_years)}y)")
    print(f"  Test:  {len(test_years)}y")

    # Model
    print("\nCreating TELEN...")
    model = create_model(config).to(device)
    print(f"  HyperNetwork: {sum(p.numel() for p in model.hypernetwork.parameters()):,} params")

    # Build graph
    print("\nBuilding concept graph...")
    train_df = df[df["year"].isin(train_years)]
    model.build_graph(train_df)
    print(f"  Graph: {model.concept_graph.num_nodes} nodes")

    # Stage 1
    print("\n" + "=" * 60)
    print("Stage 1: Contrastive Pretraining")
    print("=" * 60)
    model = contrastive_pretrain(model, train_df, config, epochs=5, batch_size=24, lr=3e-5)

    # Stage 2
    print("\n" + "=" * 60)
    print("Stage 2: Meta-Training")
    print("=" * 60)
    model = meta_train(model, articles_by_law, laws_by_year, train_years, val_years, config, epochs=50, patience=15)

    print(f"\nDone! Model saved to: {config.output_dir}/telen_best.pt")

if __name__ == "__main__":
    main()