train_st.py

# %%
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from sentence_transformers import SentenceTransformer
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, multilabel_confusion_matrix, f1_score
import matplotlib.pyplot as plt
import seaborn as sns
import math
import os
from collections import Counter

# %%
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.backends.cudnn.benchmark = True
print(f"Using device: {device}")

os.makedirs('./models', exist_ok=True)

# %%
PAD_TOKEN = '<PAD>'
UNK_TOKEN = '<UNK>'

def build_phone_vocab(phonetic_texts):
    """Build {token: idx} from all phoneme strings in the dataset."""
    counter = Counter()
    for txt in phonetic_texts:
        counter.update(str(txt).split())
    vocab = {PAD_TOKEN: 0, UNK_TOKEN: 1}
    for token, _ in counter.most_common():
        if token not in vocab:
            vocab[token] = len(vocab)
    return vocab


def tokenise_phonemes(phonetic_text, vocab, max_len=256):
    """Convert a phoneme string to a padded integer sequence."""
    tokens = str(phonetic_text).split()[:max_len]
    ids    = [vocab.get(t, vocab[UNK_TOKEN]) for t in tokens]
    # Pad / truncate
    ids   += [vocab[PAD_TOKEN]] * (max_len - len(ids))
    mask   = [1] * len(tokens) + [0] * (max_len - len(tokens))
    return ids, mask


# %%
class RhetoricalDataset(Dataset):
    """
    Each sample:
        text_emb   : (384,)  float32  — frozen sentence-transformer embedding
        phone_ids  : (256,)  int64    — tokenised phoneme sequence
        phone_mask : (256,)  float32  — 1 = real token, 0 = pad
        label      : (C,)    float32  — multi-hot label vector
    """

    PHONE_MAX_LEN = 256

    def __init__(self, df: pd.DataFrame, sent_encoder: SentenceTransformer,
                 phone_vocab: dict=None, fit_vocab: bool = True):

        for col in ['full_text', 'full_text_phonetic', 'highlights', 'highlights_phonetic']:
            df[col] = df[col].fillna("")

        self.classes      = sorted(df['figure_name'].unique())
        self.num_classes  = len(self.classes)
        self.class_to_idx = {c: i for i, c in enumerate(self.classes)}

        # Augment with highlight spans
        base = df[['full_text', 'full_text_phonetic', 'figure_name']].copy()

        span_rows = []
        for _, row in df.iterrows():
            spans_txt = [s.strip() for s in str(row['highlights']).split(';')     if s.strip()]
            spans_phn = [s.strip() for s in str(row['highlights_phonetic']).split(';') if s.strip()]
            for st, sp in zip(spans_txt, spans_phn):
                if len(st) >= 3 and st.lower() != str(row['full_text']).strip().lower():
                    span_rows.append({
                        'full_text':          st,
                        'full_text_phonetic': sp,
                        'figure_name':        row['figure_name'],
                    })

        augmented = pd.concat([base, pd.DataFrame(span_rows)], ignore_index=True)

        # Collapse multi-label annotations
        self.grouped = (
            augmented.groupby('full_text')
            .agg(
                full_text_phonetic=('full_text_phonetic', 'first'),
                figure_name       =('figure_name', lambda x: list(set(x)))
            )
            .reset_index()
        )

        print(f"Dataset: {len(self.grouped)} unique texts "
              f"(original + {len(span_rows)} highlight spans)")

        # Build / reuse phoneme vocab
        all_phonetic = self.grouped['full_text_phonetic'].tolist()
        if fit_vocab:
            self.phone_vocab = build_phone_vocab(all_phonetic)
        else:
            assert phone_vocab is not None
            self.phone_vocab = phone_vocab

        # Sentence embeddings
        print("Encoding text with sentence-transformer (batched)...")
        texts = self.grouped['full_text'].tolist()
        self.text_embeddings = sent_encoder.encode(
            texts,
            batch_size=128,
            show_progress_bar=True,
            convert_to_numpy=True,
            normalize_embeddings=True,
        ).astype(np.float32)

        # Phoneme token sequences
        print("Tokenising phoneme sequences...")
        self.phone_ids   = []
        self.phone_masks = []
        for phn in all_phonetic:
            ids, mask = tokenise_phonemes(phn, self.phone_vocab, self.PHONE_MAX_LEN)
            self.phone_ids.append(ids)
            self.phone_masks.append(mask)

        self.phone_ids   = np.array(self.phone_ids,   dtype=np.int64)
        self.phone_masks = np.array(self.phone_masks, dtype=np.float32)

        # Multi-hot labels
        self.multi_labels = []
        for figures in self.grouped['figure_name']:
            vec = np.zeros(self.num_classes, dtype=np.float32)
            for fig in figures:
                if fig in self.class_to_idx:
                    vec[self.class_to_idx[fig]] = 1.0
            self.multi_labels.append(vec)

        self.label_matrix = np.array(self.multi_labels)

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

    def __getitem__(self, idx):
        return (
            torch.tensor(self.text_embeddings[idx]),          # (384,)
            torch.tensor(self.phone_ids[idx]),                # (256,)
            torch.tensor(self.phone_masks[idx]),              # (256,)
            torch.tensor(self.multi_labels[idx]),             # (C,)
        )


# %%
class PhoneticBiLSTM(nn.Module):
    """
    Encodes a padded phoneme-token sequence into a fixed-size vector.

    Architecture:
        Embedding → 3-layer BiLSTM → masked mean pool → LayerNorm → Linear
    """

    def __init__(self, vocab_size: int, embed_dim: int = 64,
                 hidden_dim: int = 256, num_layers: int = 3,
                 out_dim: int = 384, dropout: float = 0.3):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.lstm      = nn.LSTM(
            input_size    = embed_dim,
            hidden_size   = hidden_dim,
            num_layers    = num_layers,
            batch_first   = True,
            bidirectional = True,
            dropout       = dropout if num_layers > 1 else 0.0,
        )
        self.norm    = nn.LayerNorm(hidden_dim * 2)
        self.project = nn.Linear(hidden_dim * 2, out_dim)

    def forward(self, ids: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        emb, _ = self.lstm(self.embedding(ids))  # (B, L, 2H)
        mask_3d = mask.unsqueeze(-1) 
        emb = emb.masked_fill(mask_3d == 0, -1e9)
        pooled, _ = torch.max(emb, dim=1) 
        
        return self.project(self.norm(pooled))


class RhetoricalClassifier(nn.Module):
    """
    Dual-stream classifier:
        Stream A — frozen sentence-transformer embedding (384-d)
                   captures semantic/structural meaning and POSITION
        Stream B — trained BiLSTM over phoneme tokens (384-d)
                   captures sound patterns with full sequence order

    Both streams are projected to a shared 384-d space, then combined
    via a learned per-class gate before a shared classifier head.
    """

    def __init__(self, num_classes: int, vocab_size: int, class_names: list,
                 embed_dim: int = 384, dropout: float = 0.4):
        super().__init__()

        self.phone_encoder = PhoneticBiLSTM(
            vocab_size = vocab_size,
            embed_dim  = 64,
            hidden_dim = 256,
            num_layers = 3,
            out_dim    = embed_dim,
        )

        # Per-class stream gate (text vs phonetic)
        self.stream_weights = nn.Parameter(torch.zeros(num_classes, 2))
        self._apply_gating_prior(class_names)

        # Shared classifier head
        self.head = nn.Sequential(
            nn.Linear(embed_dim, 256),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(256, 1)
        )

    def _apply_gating_prior(self, class_names):
        phonetic_prior = ['alliteration', 'assonance', 'consonance', 'rhyme']
        with torch.no_grad():
            for i, name in enumerate(class_names):
                if name.lower() in phonetic_prior:
                    # Give phonetic stream a significant head start (approx 80% weight)
                    self.stream_weights[i, 1] = 1.5 
                else:
                    # Structural figures (Antimetabole, etc) start balanced
                    self.stream_weights[i, 0] = 0.5

    def forward(self, text_emb: torch.Tensor,
                phone_ids: torch.Tensor,
                phone_mask: torch.Tensor) -> torch.Tensor:
        
        phone_emb = self.phone_encoder(phone_ids, phone_mask)
        gates = torch.softmax(self.stream_weights, dim=1) # (C, 2)

        # Multiply each stream by its respective per-class gate
        fused = (
            text_emb.unsqueeze(1) * gates[:, 0].unsqueeze(-1) + 
            phone_emb.unsqueeze(1) * gates[:, 1].unsqueeze(-1)
        )

        # Logits: (B, C)
        logits = self.head(fused).squeeze(-1)
        return logits


# %%
def compute_pos_weights(label_matrix: np.ndarray, device,
                        min_w: float = 1.0, max_w: float = 20.0) -> torch.Tensor:
    """
    pos_weight[i] = neg_count[i] / pos_count[i], clamped to [min_w, max_w].
    Passed to BCEWithLogitsLoss so rare figures receive proportionally
    larger gradient signal.
    """
    pos = label_matrix.sum(axis=0).clip(min=1)
    neg = len(label_matrix) - pos
    w   = np.clip(neg / pos, min_w, max_w)
    return torch.tensor(w, dtype=torch.float32).to(device)


def train_and_tune(dataset, train_loader, test_loader, epochs: int = 100):
    model   = RhetoricalClassifier(
        num_classes = dataset.num_classes,
        vocab_size  = len(dataset.phone_vocab),
        class_names = dataset.classes
    ).to(device)

    pos_w     = compute_pos_weights(dataset.label_matrix, device)
    criterion = nn.BCEWithLogitsLoss(pos_weight=pos_w)

    print("\nPer-class positive weights:")
    for i, cls in enumerate(dataset.classes):
        print(f"  {cls:<15} {pos_w[i].item():.2f}")

    optimizer = optim.AdamW(model.parameters(), lr=2e-4, weight_decay=1e-4)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(
        optimizer, T_max=epochs, eta_min=1e-6
    )

    print(f"\nTraining for {epochs} epochs...")
    for epoch in range(epochs):
        model.train()
        total_loss = 0.0

        for text_emb, phone_ids, phone_mask, labels in train_loader:
            text_emb   = text_emb.to(device)
            phone_ids  = phone_ids.to(device)
            phone_mask = phone_mask.to(device)
            labels     = labels.to(device)

            optimizer.zero_grad()
            loss = criterion(model(text_emb, phone_ids, phone_mask), labels)
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()
            total_loss += loss.item()

        scheduler.step()

        if (epoch + 1) % 10 == 0:
            lr = scheduler.get_last_lr()[0]
            print(f"Epoch {epoch+1:3d} | Loss: {total_loss/len(train_loader):.4f} | LR: {lr:.2e}")

    # Threshold optimisation
    print("\nOptimising decision thresholds on test set...")
    model.eval()
    all_probs, all_labels = [], []

    with torch.no_grad():
        for text_emb, phone_ids, phone_mask, labels in test_loader:
            probs = torch.sigmoid(
                model(text_emb.to(device), phone_ids.to(device), phone_mask.to(device))
            ).cpu().numpy()
            all_probs.extend(probs)
            all_labels.extend(labels.numpy())

    all_probs  = np.array(all_probs)
    all_labels = np.array(all_labels)

    best_thresholds = []
    for i, cls in enumerate(dataset.classes):
        pos_count      = int(all_labels[:, i].sum())
        n_steps        = 100 if pos_count < 50 else 50
        best_f1, best_t = 0.0, 0.5

        for t in np.linspace(0.05, 0.95, n_steps):
            preds = (all_probs[:, i] > t).astype(float)
            f1    = f1_score(all_labels[:, i], preds, zero_division=0)
            if f1 > best_f1:
                best_f1, best_t = f1, t

        best_thresholds.append(best_t)
        print(f"  {cls:<15} pos={pos_count:4d} | thresh={best_t:.3f} | F1={best_f1:.4f}")

    return model, best_thresholds, all_probs, all_labels


# %%
def plot_multilabel_cm(labels, preds, class_names):
    mcm  = multilabel_confusion_matrix(labels, preds)
    cols = 4
    rows = math.ceil(len(class_names) / cols)
    fig, axes = plt.subplots(rows, cols, figsize=(16, rows * 3))
    axes = axes.ravel()

    for i, name in enumerate(class_names):
        sns.heatmap(mcm[i], annot=True, fmt='d', ax=axes[i],
                    cmap='Greens', cbar=False,
                    xticklabels=['Pred N', 'Pred P'],
                    yticklabels=['True N', 'True P'])
        axes[i].set_title(name)

    for j in range(i + 1, len(axes)):
        axes[j].set_visible(False)

    plt.suptitle('Per-Figure Confusion Matrices', fontsize=14, y=1.01)
    plt.tight_layout()
    plt.savefig('./models/confusion_matrices.png', dpi=150, bbox_inches='tight')
    plt.show()


def plot_gate_weights(model, class_names):
    w = torch.sigmoid(model.stream_weights).detach().cpu().numpy()
    df_w = pd.DataFrame(w, columns=['Text (Transformer)', 'Phonetic (BiLSTM)'],
                        index=class_names)
    plt.figure(figsize=(8, 7))
    sns.heatmap(df_w, annot=True, fmt='.2f', cmap='YlGnBu', vmin=0, vmax=1)
    plt.title('Learned Stream Gate Weights per Rhetorical Figure')
    plt.tight_layout()
    plt.savefig('./models/gate_weights.png', dpi=150, bbox_inches='tight')
    plt.show()


def run_diagnostics(model, test_loader, best_thresholds, class_names, save_name,
                    phone_vocab):
    model.eval()
    all_probs, all_labels = [], []

    with torch.no_grad():
        for text_emb, phone_ids, phone_mask, labels in test_loader:
            probs = torch.sigmoid(
                model(text_emb.to(device), phone_ids.to(device), phone_mask.to(device))
            ).cpu().numpy()
            all_probs.extend(probs)
            all_labels.extend(labels.numpy())

    all_probs  = np.array(all_probs)
    all_labels = np.array(all_labels)
    preds      = (all_probs > np.array(best_thresholds)).astype(float)

    print("\nFinal Multi-Label Classification Report:")
    print(classification_report(all_labels, preds, target_names=class_names, digits=4))

    plot_multilabel_cm(all_labels, preds, class_names)

    # Co-occurrence
    co = np.dot(preds.T, preds)
    d  = np.diagonal(co)
    with np.errstate(divide='ignore', invalid='ignore'):
        co_norm = np.nan_to_num(co / d[:, None])

    plt.figure(figsize=(14, 12))
    sns.heatmap(co_norm, annot=True, fmt='.2f', cmap='YlGnBu',
                xticklabels=class_names, yticklabels=class_names)
    plt.title('Rhetorical Figure Co-occurrence (Normalised)')
    plt.tight_layout()
    plt.savefig('./models/cooccurrence.png', dpi=150, bbox_inches='tight')
    plt.show()

    plot_gate_weights(model, class_names)

    save_path = f'./models/{save_name}.pth'
    torch.save({
        'model_state_dict': model.state_dict(),
        'classes':          class_names,
        'thresholds':       best_thresholds,
        'phone_vocab':      phone_vocab,
    }, save_path)
    print(f"\nModel saved → {save_path}")

    return all_labels, preds


# %%
print("Loading sentence-transformer (downloads on first run ~80 MB)...")
sent_encoder = SentenceTransformer('all-MiniLM-L6-v2')
sent_encoder.to(device)

df      = pd.read_csv('./training/gofigure_phonetized.csv')
dataset = RhetoricalDataset(df, sent_encoder, fit_vocab=True)

indices             = np.arange(len(dataset))
train_idx, test_idx = train_test_split(indices, test_size=0.2, random_state=42)

train_loader = DataLoader(
    torch.utils.data.Subset(dataset, train_idx),
    batch_size=64, shuffle=True, num_workers=0, pin_memory=(device.type == 'cuda')
)
test_loader = DataLoader(
    torch.utils.data.Subset(dataset, test_idx),
    batch_size=64, shuffle=False, num_workers=0, pin_memory=(device.type == 'cuda')
)

model, thresholds, probs, labels = train_and_tune(
    dataset, train_loader, test_loader, epochs=100
)

labels, preds = run_diagnostics(
    model, test_loader, thresholds, dataset.classes,
    save_name='rhetoric_multilabel_v4',
    phone_vocab=dataset.phone_vocab,
)


# %%
def identify_figures(text: str, phonetic_text: str,
                     model, checkpoint, sent_encoder):
    """
    Run inference on a single example and print detected figures.
    """
    phone_vocab = checkpoint['phone_vocab']

    # Text embedding
    text_emb = sent_encoder.encode(
        [text], convert_to_numpy=True, normalize_embeddings=True
    ).astype(np.float32)                                    # (1, 384)

    # Phoneme tokenisation
    ids, mask = tokenise_phonemes(phonetic_text, phone_vocab, max_len=256)

    text_tensor  = torch.tensor(text_emb).to(device)
    phone_tensor = torch.tensor([ids], dtype=torch.long).to(device)
    mask_tensor  = torch.tensor([mask], dtype=torch.float32).to(device)

    with torch.no_grad():
        scores = torch.sigmoid(
            model(text_tensor, phone_tensor, mask_tensor)
        ).cpu().numpy()[0]

    print(f'\n{"="*60}')
    print(f'Input: "{text}"')
    print(f'{"="*60}')

    detected = False
    for i, score in enumerate(scores):
        thresh = checkpoint['thresholds'][i]
        if score > thresh:
            detected = True
            fig = checkpoint['classes'][i].upper()
            print(f'  [{fig}]  confidence={score:.3f}  threshold={thresh:.3f}')

    if not detected:
        print('  No figures detected above threshold.')
    print()


# Load checkpoint and run examples
checkpoint = torch.load(
    './models/rhetoric_multilabel_v4.pth',
    map_location=device, weights_only=False
)
inference_model = RhetoricalClassifier(
    num_classes = len(checkpoint['classes']),
    vocab_size  = len(checkpoint['phone_vocab']),
    class_names = checkpoint['classes']
).to(device)
inference_model.load_state_dict(checkpoint['model_state_dict'])
inference_model.eval()

identify_figures(
    "The general who became a slave; the slave who became a gladiator; the gladiator who defied an Emperor",
    "ð ə <W> dʒ ɛ n ɹ ə l <W> h uː <W> b ɪ k eɪ m <W> ə <W> s l eɪ v <W> ð ə <W> s l eɪ v <W> h uː <W> b ɪ k eɪ m <W> ə <W> ɡ l æ d i eɪ t ə ɹ <W> ð ə <W> ɡ l æ d i eɪ t ə ɹ <W> h uː <W> d ɪ f aɪ d <W> æ n <W> ɛ m p ə ɹ ə ɹ",
    inference_model, checkpoint, sent_encoder
)

identify_figures(
    "The rain in Spain stays mainly in the plain",
    "ð ə <W> ɹ eɪ n <W> ɪ n <W> s p eɪ n <W> s t eɪ z <W> m eɪ n l i <W> ɪ n <W> ð ə <W> p l eɪ n",
    inference_model, checkpoint, sent_encoder
)

identify_figures(
    "Peter Piper picked a peck of pickled peppers",
    "p iː t ə ɹ <W> p aɪ p ə ɹ <W> p ɪ k t <W> ɐ <W> p ɛ k <W> ʌ v <W> p ɪ k l d <W> p ɛ p ə ɹ z",
    inference_model, checkpoint, sent_encoder
)

identify_figures(
    "The odds are good, but the goods are odd.",
    "ð ɪ <W> ɑː d z <W> ɑːɹ <W> ɡ ʊ d <W> b ʌ t <W> ð ə <W> ɡ ʊ d z <W> ɑː ɹ <W> ɑː d",
    inference_model, checkpoint, sent_encoder
)