downstream_eval.py

"""
Downstream evaluation for ModernProteinLM on predictive protein tasks:
- Fluorescence (regression, Spearman)
- Solubility (binary classification)
- Secondary Structure (token classification, Q3/Q8 accuracy)
- Remote Homology (classification)

Compares against ESM-2 baselines.
"""

import os
import json
import torch
import torch.nn as nn
import numpy as np
from torch.utils.data import DataLoader, Dataset
from sklearn.metrics import accuracy_score, f1_score, roc_auc_score, mean_squared_error
from scipy.stats import spearmanr
from transformers import get_linear_schedule_with_warmup
from datasets import load_dataset
from tqdm import tqdm
import warnings
warnings.filterwarnings("ignore")

from modeling_modern_protein import ModernProteinLM, ModernProteinLMConfig
from electra_pretrain import ProteinTokenizer


class ProteinDownstreamDataset(Dataset):
    """Generic downstream dataset wrapper."""
    
    TASK_CONFIGS = {
        "fluorescence": {
            "dataset": "proteinea/fluorescence",
            "seq_col": "primary",
            "label_col": "log_fluorescence",
            "task": "regression",
            "metric": "spearman",
        },
        "solubility": {
            "dataset": "proteinea/solubility",
            "seq_col": "sequences",
            "label_col": "labels",
            "task": "classification",
            "num_labels": 2,
            "metric": "accuracy",
        },
        "secondary_structure": {
            "dataset": "proteinea/secondary_structure_prediction",
            "seq_col": "input",
            "label_cols": ["dssp3", "dssp8"],
            "task": "token_classification",
            "num_labels": 3,  # Q3 first
            "metric": "accuracy",
        },
        "remote_homology": {
            "dataset": "proteinea/remote_homology",
            "seq_col": "primary",
            "label_col": "fold_label",
            "task": "classification",
            "num_labels": 1195,  # Actually fold labels
            "metric": "accuracy",
        },
    }
    
    def __init__(self, task_name, split, tokenizer, max_length=1024):
        self.task_name = task_name
        self.config = self.TASK_CONFIGS[task_name]
        self.tokenizer = tokenizer
        self.max_length = max_length
        
        try:
            self.data = load_dataset(self.config["dataset"], split=split)
        except:
            # Some datasets don't have validation/test splits, use train
            self.data = load_dataset(self.config["dataset"], split="train")
        
        self.examples = list(self.data)
    
    def __len__(self):
        return len(self.examples)
    
    def __getitem__(self, idx):
        ex = self.examples[idx]
        seq = ex[self.config["seq_col"]]
        encoded = self.tokenizer.encode(seq, max_length=self.max_length)
        
        item = {
            "input_ids": torch.tensor(encoded["input_ids"], dtype=torch.long),
            "attention_mask": torch.tensor(encoded["attention_mask"], dtype=torch.long),
        }
        
        if self.config["task"] == "regression":
            item["labels"] = torch.tensor(ex[self.config["label_col"]], dtype=torch.float)
        elif self.config["task"] == "classification":
            item["labels"] = torch.tensor(ex[self.config["label_col"]], dtype=torch.long)
        elif self.config["task"] == "token_classification":
            # Secondary structure: each AA has a label
            ss = ex[self.config["label_cols"][0]]  # dssp3
            # Map 'C', 'H', 'E' to 0, 1, 2
            ss_map = {'C': 0, 'H': 1, 'E': 2}
            labels = [ss_map.get(c, 0) for c in ss]
            # Pad/truncate to match sequence length
            seq_len = sum(encoded["attention_mask"])
            labels = labels[:seq_len]
            while len(labels) < len(encoded["input_ids"]):
                labels.append(-100)
            item["labels"] = torch.tensor(labels, dtype=torch.long)
        
        return item


class DownstreamModel(nn.Module):
    def __init__(self, base_model, task_config):
        super().__init__()
        self.base = base_model
        self.task = task_config["task"]
        self.config = task_config
        
        hidden_size = base_model.config.hidden_size
        
        if self.task == "regression":
            self.head = nn.Linear(hidden_size, 1)
        elif self.task == "classification":
            self.head = nn.Linear(hidden_size, task_config.get("num_labels", 2))
        elif self.task == "token_classification":
            self.head = nn.Linear(hidden_size, task_config.get("num_labels", 3))
    
    def forward(self, input_ids, attention_mask, labels=None):
        outputs = self.base(
            input_ids=input_ids,
            attention_mask=attention_mask,
            output_hidden_states=True,
            return_dict=True,
        )
        hidden = outputs.hidden_states[-1]
        
        if self.task in ["regression", "classification"]:
            # Mean pool
            mask_expanded = attention_mask.unsqueeze(-1).float()
            pooled = (hidden * mask_expanded).sum(dim=1) / mask_expanded.sum(dim=1).clamp(min=1e-9)
            logits = self.head(pooled)
        else:
            # Token-level
            logits = self.head(hidden)
        
        loss = None
        if labels is not None:
            if self.task == "regression":
                loss_fct = nn.MSELoss()
                loss = loss_fct(logits.squeeze(-1), labels)
            elif self.task == "classification":
                loss_fct = nn.CrossEntropyLoss()
                loss = loss_fct(logits, labels)
            elif self.task == "token_classification":
                loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
                loss = loss_fct(logits.view(-1, self.config.get("num_labels", 3)), labels.view(-1))
        
        return {"loss": loss, "logits": logits}


def evaluate(model, dataloader, task_config, device):
    model.eval()
    all_preds = []
    all_labels = []
    total_loss = 0
    
    with torch.no_grad():
        for batch in dataloader:
            input_ids = batch["input_ids"].to(device)
            attention_mask = batch["attention_mask"].to(device)
            labels = batch["labels"].to(device)
            
            outputs = model(input_ids, attention_mask, labels)
            total_loss += outputs["loss"].item() * input_ids.size(0)
            
            logits = outputs["logits"]
            if task_config["task"] == "regression":
                preds = logits.squeeze(-1).cpu().numpy()
                all_preds.extend(preds)
                all_labels.extend(labels.cpu().numpy())
            elif task_config["task"] == "classification":
                preds = torch.argmax(logits, dim=-1).cpu().numpy()
                all_preds.extend(preds)
                all_labels.extend(labels.cpu().numpy())
            elif task_config["task"] == "token_classification":
                preds = torch.argmax(logits, dim=-1).cpu().numpy()
                labels_np = labels.cpu().numpy()
                # Only evaluate non-padding positions
                for i in range(len(preds)):
                    mask = labels_np[i] != -100
                    all_preds.extend(preds[i][mask])
                    all_labels.extend(labels_np[i][mask])
    
    metric = task_config["metric"]
    if metric == "spearman":
        score, _ = spearmanr(all_labels, all_preds)
    elif metric == "accuracy":
        score = accuracy_score(all_labels, all_preds)
    elif metric == "f1":
        score = f1_score(all_labels, all_preds, average="macro")
    
    avg_loss = total_loss / len(dataloader.dataset)
    return score, avg_loss


def train_downstream(
    base_model,
    task_name,
    tokenizer,
    epochs=20,
    batch_size=16,
    lr=1e-4,
    device="cuda",
    seed=42,
):
    torch.manual_seed(seed)
    np.random.seed(seed)
    
    task_config = ProteinDownstreamDataset.TASK_CONFIGS[task_name]
    
    train_dataset = ProteinDownstreamDataset(task_name, "train", tokenizer)
    
    # For validation, use test or create split
    try:
        val_dataset = ProteinDownstreamDataset(task_name, "validation", tokenizer)
    except:
        val_dataset = ProteinDownstreamDataset(task_name, "test", tokenizer)
    
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=2)
    
    model = DownstreamModel(base_model, task_config).to(device)
    
    # Freeze some layers for small datasets
    if task_name in ["fluorescence"]:
        # Fine-tune all for small regression tasks
        pass
    
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
    
    total_steps = len(train_loader) * epochs
    scheduler = get_linear_schedule_with_warmup(
        optimizer, num_warmup_steps=int(0.1 * total_steps), num_training_steps=total_steps
    )
    
    best_score = -float("inf") if task_config["metric"] != "mse" else float("inf")
    best_model_state = None
    
    for epoch in range(epochs):
        model.train()
        total_loss = 0
        
        pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}")
        for batch in pbar:
            input_ids = batch["input_ids"].to(device)
            attention_mask = batch["attention_mask"].to(device)
            labels = batch["labels"].to(device)
            
            outputs = model(input_ids, attention_mask, labels)
            loss = outputs["loss"]
            
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            scheduler.step()
            optimizer.zero_grad()
            
            total_loss += loss.item()
            pbar.set_postfix({"loss": f"{loss.item():.4f}"})
        
        # Evaluate
        score, val_loss = evaluate(model, val_loader, task_config, device)
        print(f"Epoch {epoch+1}: Val {task_config['metric']}={score:.4f}, Loss={val_loss:.4f}")
        
        if task_config["metric"] == "spearman":
            is_better = score > best_score
        elif task_config["metric"] == "accuracy":
            is_better = score > best_score
        
        if is_better:
            best_score = score
            best_model_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
    
    if best_model_state:
        model.load_state_dict(best_model_state)
    
    return model, best_score


def compare_models(
    task_names=["fluorescence", "solubility", "secondary_structure"],
    epochs=20,
    device="cuda",
):
    tokenizer = ProteinTokenizer()
    results = {}
    
    for task in task_names:
        print(f"\n{'='*50}")
        print(f"Task: {task}")
        print(f"{'='*50}")
        
        # ModernProteinLM (random init)
        config = ModernProteinLMConfig(
            vocab_size=33,
            hidden_size=640,
            num_hidden_layers=24,
            num_attention_heads=10,
            intermediate_size=2304,
            use_geglu=True,
            tie_word_embeddings=True,
        )
        modern_model = ModernProteinLM(config)
        print(f"ModernProteinLM params: {sum(p.numel() for p in modern_model.parameters())/1e6:.1f}M")
        
        modern_model, modern_score = train_downstream(
            modern_model, task, tokenizer, epochs=epochs, device=device
        )
        
        # ESM-2 baseline
        try:
            from transformers import AutoModel, AutoTokenizer
            esm_tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t12_35M_UR50D")
            esm_model = AutoModel.from_pretrained("facebook/esm2_t12_35M_UR50D")
            print(f"ESM-2 35M params: {sum(p.numel() for p in esm_model.parameters())/1e6:.1f}M")
            
            # Convert ESM model to have same interface
            esm_model.config.hidden_size = esm_model.config.hidden_size
            
            esm_model, esm_score = train_downstream(
                esm_model, task, tokenizer, epochs=epochs, device=device
            )
            
            results[task] = {
                "modern": modern_score,
                "esm2_35m": esm_score,
            }
        except Exception as e:
            print(f"ESM-2 comparison failed: {e}")
            results[task] = {"modern": modern_score, "esm2_35m": None}
        
        print(f"\nResults for {task}:")
        print(f"  ModernProteinLM: {modern_score:.4f}")
        if "esm2_35m" in results[task] and results[task]["esm2_35m"] is not None:
            print(f"  ESM-2 35M: {results[task]['esm2_35m']:.4f}")
    
    with open("downstream_results.json", "w") as f:
        json.dump(results, f, indent=2)
    
    return results


if __name__ == "__main__":
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Using device: {device}")
    
    # Quick test on solubility (smallest dataset)
    tokenizer = ProteinTokenizer()
    
    config = ModernProteinLMConfig(
        vocab_size=33,
        hidden_size=128,
        num_hidden_layers=4,
        num_attention_heads=4,
        intermediate_size=512,
        use_geglu=True,
        tie_word_embeddings=True,
    )
    model = ModernProteinLM(config)
    
    print(f"\nTesting on solubility (tiny model)...")
    trained_model, score = train_downstream(
        model, "solubility", tokenizer, epochs=5, batch_size=8, lr=5e-4, device=device
    )
    print(f"Solubility accuracy: {score:.4f}")