app.py

import spaces
import os
import json
import random
import math
import logging
import traceback
from pathlib import Path
from dataclasses import dataclass
from typing import Dict, List, Tuple, Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForCausalLM

import gradio as gr
import pandas as pd
import plotly.graph_objects as go

# ---------------------------------------------------------------------------
# Logging
# ---------------------------------------------------------------------------
logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s")
logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Constants & Paths
# ---------------------------------------------------------------------------
MODEL_IDS: List[str] = [
    "CodeSoft/MetaDiffusion-150M-ChatBase",
    "BananaMind/BananaMind-2-Medium-Chat",
    "SupraLabs/Supra2-100M-Instruct",
    "HuggingFaceTB/SmolLM2-135M-Instruct",
    "OpenCerebral/Boris-1.3-125M-Instruct",
    "OpenCerebral/Boris-1.3-75M-Instruct",
]

MODEL_DISPLAY: Dict[str, str] = {
    "CodeSoft/MetaDiffusion-150M-ChatBase": "MetaDiffusion-150M-ChatBase",
    "BananaMind/BananaMind-2-Medium-Chat": "BananaMind-2-Medium-Chat",
    "SupraLabs/Supra2-100M-Instruct": "Supra2-100M-Instruct",
    "HuggingFaceTB/SmolLM2-135M-Instruct": "SmolLM2-135M-Instruct",
    "OpenCerebral/Boris-1.3-125M-Instruct": "Boris-1.3-125M-Instruct",
    "OpenCerebral/Boris-1.3-75M-Instruct": "Boris-1.3-75M-Instruct",
}

BASE_MODEL_IDS: List[str] = [
    "fromziro/Zero-v0.1-150M",
    "AxiomicLabs/GPT-X2.5-135M",
    "BananaMind/BananaMind-2-Pro",
    "HuggingFaceTB/SmolLM2-135M",
    "OpenCerebral/Boris-1.3-125M",
    "OpenCerebral/Boris-1.3-75M",
]

BASE_MODEL_DISPLAY: Dict[str, str] = {
    "fromziro/Zero-v0.1-150M": "Zero-v0.1-150M",
    "AxiomicLabs/GPT-X2.5-135M": "GPT-X2.5-135M",
    "BananaMind/BananaMind-2-Pro": "BananaMind-2-Pro",
    "HuggingFaceTB/SmolLM2-135M": "SmolLM2-135M",
    "OpenCerebral/Boris-1.3-125M": "Boris-1.3-125M",
    "OpenCerebral/Boris-1.3-75M": "Boris-1.3-75M",
}

MODEL_PARAMS: Dict[str, float] = {
    "CodeSoft/MetaDiffusion-150M-ChatBase": 169.5e6,
    "BananaMind/BananaMind-2-Medium-Chat": 49.6e6,
    "SupraLabs/Supra2-100M-Instruct": 100.0e6,
    "HuggingFaceTB/SmolLM2-135M-Instruct": 135.0e6,
    "OpenCerebral/Boris-1.3-125M-Instruct": 125.0e6,
    "OpenCerebral/Boris-1.3-75M-Instruct": 75.0e6,
    "fromziro/Zero-v0.1-150M": 151.6e6,
    "AxiomicLabs/GPT-X2.5-135M": 135.0e6,
    "BananaMind/BananaMind-2-Pro": 139.0e6,
    "HuggingFaceTB/SmolLM2-135M": 135.0e6,
    "OpenCerebral/Boris-1.3-125M": 125.0e6,
    "OpenCerebral/Boris-1.3-75M": 75.0e6,
}

FALLBACK_IDS: Dict[str, str] = {}

MODERATION_MODEL_ID = "ifmain/ModerationBERT-En-02"
MODERATION_CATEGORIES = [
    "harassment",
    "harassment_threatening",
    "hate",
    "hate_threatening",
    "self_harm",
    "self_harm_instructions",
    "self_harm_intent",
    "sexual",
    "sexual_minors",
    "violence",
    "violence_graphic",
    "self-harm",
    "sexual/minors",
    "hate/threatening",
    "violence/graphic",
    "self-harm/intent",
    "self-harm/instructions",
    "harassment/threatening",
]
MODERATION_THRESHOLD = 0.35

INIT_RATING = 1000
K_FACTOR = 32
SCALE = 400
BASE = 10

# All data in ./data
try:
    BASE_DIR = Path(__file__).parent
except NameError:
    BASE_DIR = Path(".")

# Prefer /data (HF Space bucket mount) if available, otherwise fallback to ./data
# Bucket is mounted at /data in Space — use dynamic check each call so late mounts are detected
def get_data_dir() -> Path:
    bucket = Path("/data")
    if bucket.exists() and bucket.is_dir():
        try:
            # Ensure writable (touch test)
            (bucket / ".write_test").touch(exist_ok=True)
            (bucket / ".write_test").unlink(missing_ok=True)
            return bucket
        except Exception:
            pass
    # Fallback to local ./data
    local = BASE_DIR / "data"
    try:
        local.mkdir(parents=True, exist_ok=True)
    except Exception:
        pass
    return local

def get_elo_file(name: str = "elo") -> Path:
    return get_data_dir() / f"{name}.json"

def get_chat_file(name: str = "chats") -> Path:
    return get_data_dir() / f"{name}.jsonl"

# Keep legacy globals for backwards compat (now dynamic via functions)
DATA_DIR = get_data_dir()
ELO_FILE = get_elo_file()
CHAT_FILE = get_chat_file()

GEN_DEFAULTS: Dict[str, dict] = {
    "HuggingFaceTB/SmolLM2-135M-Instruct": {"max_new_tokens": 64, "temperature": 0.7, "top_p": 0.9, "repetition_penalty": 1.1, "do_sample": True},
    "SupraLabs/Supra2-100M-Instruct": {"max_new_tokens": 64, "temperature": 0.7, "top_p": 0.9, "top_k": 25, "repetition_penalty": 1.1, "do_sample": True, "no_repeat_ngram_size": 3},
    "BananaMind/BananaMind-2-Medium-Chat": {"max_new_tokens": 64, "temperature": 0.7, "top_p": 0.9, "repetition_penalty": 1.1, "do_sample": True},
    "CodeSoft/MetaDiffusion-150M-ChatBase": {"max_new_tokens": 96, "num_steps": 128, "temperature": 0.7, "top_p": 0.9, "repetition_penalty": 1.5},
    "fromziro/Zero-v0.1-150M": {"max_new_tokens": 64, "temperature": 0.8, "top_p": 0.95, "repetition_penalty": 1.1, "do_sample": True},
    "AxiomicLabs/GPT-X2.5-135M": {"max_new_tokens": 64, "temperature": 0.8, "top_p": 0.95, "repetition_penalty": 1.1, "do_sample": True},
    "BananaMind/BananaMind-2-Pro": {"max_new_tokens": 64, "temperature": 0.8, "top_p": 0.95, "repetition_penalty": 1.1, "do_sample": True},
    "HuggingFaceTB/SmolLM2-135M": {"max_new_tokens": 64, "temperature": 0.8, "top_p": 0.95, "repetition_penalty": 1.1, "do_sample": True},
    "OpenCerebral/Boris-1.3-125M-Instruct": {"max_new_tokens": 64, "temperature": 0.7, "top_p": 0.9, "repetition_penalty": 1.1, "do_sample": True},
    "OpenCerebral/Boris-1.3-75M-Instruct": {"max_new_tokens": 64, "temperature": 0.7, "top_p": 0.9, "repetition_penalty": 1.1, "do_sample": True},
    "OpenCerebral/Boris-1.3-125M": {"max_new_tokens": 64, "temperature": 0.8, "top_p": 0.95, "repetition_penalty": 1.1, "do_sample": True},
    "OpenCerebral/Boris-1.3-75M": {"max_new_tokens": 64, "temperature": 0.8, "top_p": 0.95, "repetition_penalty": 1.1, "do_sample": True},
}

MODEL_CONTEXT: Dict[str, int] = {
    "HuggingFaceTB/SmolLM2-135M-Instruct": 2048,
    "SupraLabs/Supra2-100M-Instruct": 1024,
    "BananaMind/BananaMind-2-Medium-Chat": 3072,
    "CodeSoft/MetaDiffusion-150M-ChatBase": 5120,
    "fromziro/Zero-v0.1-150M": 2048,
    "AxiomicLabs/GPT-X2.5-135M": 2048,
    "BananaMind/BananaMind-2-Pro": 3072,
    "HuggingFaceTB/SmolLM2-135M": 2048,
    "OpenCerebral/Boris-1.3-125M-Instruct": 2048,
    "OpenCerebral/Boris-1.3-75M-Instruct": 2048,
    "OpenCerebral/Boris-1.3-125M": 2048,
    "OpenCerebral/Boris-1.3-75M": 2048,
}


@dataclass
class ArenaSpec:
    key: str
    model_ids: List[str]
    display: Dict[str, str]
    elo_name: str
    chat_name: str
    arena_title: str
    lb_title: str


MAIN_ARENA = ArenaSpec("main", MODEL_IDS, MODEL_DISPLAY, "elo", "chats", "Arena", "Leaderboard")
BASE_ARENA = ArenaSpec("base", BASE_MODEL_IDS, BASE_MODEL_DISPLAY, "base_elo", "base_chats", "Base Arena", "Base Leaderboard")

# ZeroGPU: CUDA is emulated at startup so models load onto cuda at module level;
# real GPU is only mounted inside @spaces.GPU-decorated calls.
DEVICE = os.environ.get("SLM_ARENA_DEVICE", "") or ("cuda" if torch.cuda.is_available() else "cpu")

@dataclass
class MetaDiffusionConfig:
    hidden_size: int = 768
    intermediate_size: int = 2112
    num_hidden_layers: int = 16
    num_attention_heads: int = 12
    num_key_value_heads: int = 6
    head_dim: int = 64
    vocab_size: int = 32000
    mask_vocab_size: int = 32010
    max_position_embeddings: int = 5120
    rope_theta: float = 10000.0
    rms_norm_eps: float = 1e-6
    hidden_act: str = "silu"
    timestep_emb_hidden: int = 768
    mask_token_id: int = 32000
    pad_token_id: int = 1
    mask_ratio_min: float = 0.0
    mask_ratio_max: float = 1.0
    dtype: torch.dtype = torch.float32  # type: ignore
    tie_word_embeddings: bool = False


class _RotaryEmbedding(nn.Module):
    def __init__(self, dim, max_position_embeddings=5120, base=10000.0, device=None):
        super().__init__()
        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, device=device).float() / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    @torch.no_grad()
    def forward(self, x, position_ids):
        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
        position_ids_expanded = position_ids[:, None, :].float()
        freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
        emb = torch.cat((freqs, freqs), dim=-1)
        cos = emb.cos()
        sin = emb.sin()
        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)


def _rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


def _apply_rotary_pos_emb(q, k, cos, sin):
    cos = cos.unsqueeze(1)
    sin = sin.unsqueeze(1)
    q_embed = (q * cos) + (_rotate_half(q) * sin)
    k_embed = (k * cos) + (_rotate_half(k) * sin)
    return q_embed, k_embed


class _TimestepEmbedding(nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        self.hidden_size = hidden_size
        self.mlp = nn.Sequential(
            nn.Linear(hidden_size, hidden_size * 4),
            nn.SiLU(),
            nn.Linear(hidden_size * 4, hidden_size),
        )

    def forward(self, t):
        half_dim = self.hidden_size // 2
        emb = math.log(10000.0) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, device=t.device, dtype=torch.float32) * -emb)
        emb = t[:, None].float() * emb[None, :]
        emb = torch.cat([emb.sin(), emb.cos()], dim=-1)
        return self.mlp(emb).to(t.dtype)


class _TimestepResidual(nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        self.proj = nn.Linear(hidden_size, hidden_size)
        nn.init.zeros_(self.proj.weight)
        nn.init.zeros_(self.proj.bias)

    def forward(self, x, emb):
        return x + self.proj(emb)[:, None, :]


class _RMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = eps

    def forward(self, x):
        var = x.pow(2).mean(-1, keepdim=True)
        x = x * torch.rsqrt(var + self.eps)
        return self.weight * x


class _SelfAttention(nn.Module):
    def __init__(self, config: MetaDiffusionConfig):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.head_dim = config.head_dim
        self.num_kv_groups = self.num_heads // self.num_kv_heads
        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * config.head_dim, bias=False)
        self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * config.head_dim, bias=False)
        self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * config.head_dim, bias=False)
        self.o_proj = nn.Linear(self.num_heads * config.head_dim, config.hidden_size, bias=False)
        self.rotary_emb = _RotaryEmbedding(config.head_dim, max_position_embeddings=config.max_position_embeddings, base=config.rope_theta)

    def forward(self, x, attention_mask=None, position_ids=None):
        batch, seq, _ = x.shape
        q = self.q_proj(x).view(batch, seq, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(batch, seq, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(batch, seq, self.num_kv_heads, self.head_dim).transpose(1, 2)
        cos, sin = self.rotary_emb(x, position_ids)
        q, k = _apply_rotary_pos_emb(q, k, cos, sin)
        if self.num_kv_groups > 1:
            k = k.repeat_interleave(self.num_kv_groups, dim=1)
            v = v.repeat_interleave(self.num_kv_groups, dim=1)
        out = F.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask)
        out = out.transpose(1, 2).contiguous().view(batch, seq, -1)
        return self.o_proj(out)


class _MLP(nn.Module):
    def __init__(self, config: MetaDiffusionConfig):
        super().__init__()
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)

    def forward(self, x):
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class _TransformerBlock(nn.Module):
    def __init__(self, config: MetaDiffusionConfig):
        super().__init__()
        self.input_layernorm = _RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.self_attn = _SelfAttention(config)
        self.post_attention_layernorm = _RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.mlp = _MLP(config)
        self.timestep_residual = _TimestepResidual(config.hidden_size)

    def forward(self, x, timestep_emb, attention_mask=None, position_ids=None):
        residual = x
        x = self.input_layernorm(x)
        x = self.self_attn(x, attention_mask, position_ids)
        x = residual + x
        x = self.timestep_residual(x, timestep_emb)
        residual = x
        x = self.post_attention_layernorm(x)
        x = self.mlp(x)
        x = residual + x
        x = self.timestep_residual(x, timestep_emb)
        return x


class MetaDiffusionLM(nn.Module):
    def __init__(self, config: MetaDiffusionConfig):
        super().__init__()
        self.config = config
        self.embed_tokens = nn.Embedding(config.mask_vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
        self.timestep_emb = _TimestepEmbedding(config.timestep_emb_hidden)
        self.layers = nn.ModuleList([_TransformerBlock(config) for _ in range(config.num_hidden_layers)])
        self.norm = _RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        if config.tie_word_embeddings:
            self.lm_head = None  # type: ignore
        else:
            self.lm_head = nn.Linear(config.hidden_size, config.mask_vocab_size, bias=False)
        if self.lm_head is not None:
            nn.init.normal_(self.lm_head.weight, std=0.02)

    def forward(self, input_ids, timesteps, attention_mask=None):
        batch, seq = input_ids.shape
        position_ids = torch.arange(seq, device=input_ids.device).unsqueeze(0).expand(batch, -1)
        x = self.embed_tokens(input_ids)
        t_emb = self.timestep_emb(timesteps)
        attn_mask = None
        if attention_mask is not None:
            attn_mask = ((1.0 - attention_mask[:, None, None, :].float()) * -1e9).to(x.dtype)
        for layer in self.layers:
            x = layer(x, t_emb, attn_mask, position_ids)
        x = self.norm(x)
        if self.lm_head is not None:
            logits = self.lm_head(x)
        else:
            logits = F.linear(x, self.embed_tokens.weight)
        return logits

DIFF_MASK_ID = 32000
DIFF_CHAT_TOKENS = ["<|im_start|>", "<|im_end|>"] + [f"<|r{i}|>" for i in range(1, 8)]
DIFF_IM_START, DIFF_IM_END = "<|im_start|>", "<|im_end|>"


def _ensure_diff_chat_tokens(tokenizer):
    """Add ChatML + rainbow tokens if missing (base tokenizer case). Mirrors chat.py."""
    if tokenizer.convert_tokens_to_ids(DIFF_IM_START) == tokenizer.unk_token_id:
        if len(tokenizer) == 32000:
            tokenizer.add_special_tokens({"additional_special_tokens": ["<|reserved|>"]})
        tokenizer.add_special_tokens({"additional_special_tokens": DIFF_CHAT_TOKENS})
        assert tokenizer.convert_tokens_to_ids(DIFF_IM_END) == 32002, "chat token ids wrong (collide with mask id 32000)"
    return tokenizer


def _format_diff_messages(messages):
    parts = []
    for m in messages:
        parts.append(f"{DIFF_IM_START}{m['role']}\n{m['content']}{DIFF_IM_END}")
    return "\n".join(parts)


def _diff_cumulative_unmask_frac(i, N):
    return 0.5 * (1 - math.cos(math.pi * i / N))


def _diff_cut_response(tokens, tokenizer):
    """Cut at <|im_end|> or </s>; drop rainbow/pad. Mirrors chat.py."""
    im_end_id = tokenizer.convert_tokens_to_ids(DIFF_IM_END)
    eos_id = tokenizer.eos_token_id
    rainbow_ids = {tokenizer.convert_tokens_to_ids(f"<|r{i}|>") for i in range(1, 8)}
    out = []
    for t in tokens:
        if t == im_end_id or t == eos_id:
            break
        if t in rainbow_ids or t == tokenizer.pad_token_id:
            continue
        out.append(t)
    return out


@torch.no_grad()
def _diff_generate_response(model, tokenizer, prompt_ids, gen_len, num_steps, temperature, repetition_penalty, device, stop_on_end=True):
    model.eval()
    total_len = prompt_ids.shape[1] + gen_len
    x = torch.full((1, total_len), DIFF_MASK_ID, device=device, dtype=torch.long)
    x[0, : prompt_ids.shape[1]] = prompt_ids
    mask_id = DIFF_MASK_ID
    im_end_id = tokenizer.convert_tokens_to_ids(DIFF_IM_END)
    eos_id = tokenizer.eos_token_id
    prompt_len = prompt_ids.shape[1]

    for i in range(num_steps):
        frac_now = _diff_cumulative_unmask_frac(i, num_steps)
        frac_next = _diff_cumulative_unmask_frac(i + 1, num_steps)
        n_masked = (x == mask_id).sum().item()
        n_total = int((frac_next - frac_now) * gen_len + 0.5)
        if i == num_steps - 1:
            n_unmask = n_masked
        else:
            n_unmask = max(n_total, 1) if n_masked > 0 else 0

        t = 1.0 - frac_now
        logits = model(x, torch.full((1,), t, device=device))
        logits[:, :, mask_id] = -1e9

        if repetition_penalty != 1.0:
            for tok in x[0].unique():
                ti = int(tok.item())
                if 0 <= ti < logits.shape[-1]:
                    logits[0, :, ti] = torch.where(
                        logits[0, :, ti] < 0,
                        logits[0, :, ti] * repetition_penalty,
                        logits[0, :, ti] / repetition_penalty,
                    )

        mask_positions = x == mask_id
        if not mask_positions.any():
            break
        mask_logits = logits[mask_positions]
        probs = F.softmax(mask_logits / max(0.1, temperature), dim=-1)
        sampled = torch.multinomial(probs, 1).squeeze(-1)
        mask_flat = mask_positions.nonzero(as_tuple=False)

        if n_unmask < int(mask_positions.sum().item()):
            fill_positions = mask_flat[:n_unmask]
            for idx, tok in zip(fill_positions, sampled[:n_unmask]):
                x[idx[0], idx[1]] = tok
        else:
            x[mask_positions] = sampled

        if stop_on_end and ((x[0, prompt_len:] == im_end_id).any() or (x[0, prompt_len:] == eos_id).any()):
            break
    return x


# ---------------------------------------------------------------------------
# ELO persistence
# ---------------------------------------------------------------------------
def init_elo_state(spec: ArenaSpec) -> Dict[str, dict]:
    return {mid: {"rating": float(INIT_RATING), "wins": 0, "losses": 0, "battles": 0, "ties": 0, "both_bad": 0} for mid in spec.model_ids}

def load_elo(spec: ArenaSpec) -> Dict[str, dict]:
    if get_elo_file(spec.elo_name).exists():
        try:
            with open(get_elo_file(spec.elo_name), "r") as f:
                data = json.load(f)
            for mid in spec.model_ids:
                if mid not in data:
                    data[mid] = {"rating": float(INIT_RATING), "wins": 0, "losses": 0, "battles": 0, "ties": 0, "both_bad": 0}
                else:
                    data[mid].setdefault("rating", float(INIT_RATING))
                    data[mid].setdefault("wins", 0)
                    data[mid].setdefault("losses", 0)
                    data[mid].setdefault("battles", 0)
                    data[mid].setdefault("ties", 0)
                    data[mid].setdefault("both_bad", 0)
            return data
        except Exception as e:
            logger.warning(f"Failed to load ELO file: {e}, resetting")
    return init_elo_state(spec)

def save_elo(state: Dict[str, dict], spec: ArenaSpec):
    try:
        get_data_dir().mkdir(parents=True, exist_ok=True)
        with open(get_elo_file(spec.elo_name), "w") as f:
            json.dump(state, f, indent=2)
    except Exception as e:
        logger.error(f"Failed to save ELO: {e}")

def expected_score(ra: float, rb: float) -> float:
    return 1.0 / (1.0 + BASE ** ((rb - ra) / SCALE))

def update_elo(state: Dict[str, dict], model_a: str, model_b: str, winner: Optional[str], spec: ArenaSpec) -> Dict[str, dict]:
    if model_a not in state or model_b not in state:
        logger.warning(f"Unknown models in ELO update: {model_a}, {model_b}")
        return state
    ra = state[model_a]["rating"]
    rb = state[model_b]["rating"]
    ea = expected_score(ra, rb)
    eb = expected_score(rb, ra)
    if winner == model_a:
        sa = 1.0
    elif winner == model_b:
        sa = 0.0
    elif winner is None or winner == "tie" or winner == "both_bad":
        sa = 0.5
    else:
        raise ValueError(f"Unexpected winner: {winner}")
    sb = 1.0 - sa
    state[model_a]["rating"] = ra + K_FACTOR * (sa - ea)
    state[model_b]["rating"] = rb + K_FACTOR * (sb - eb)
    state[model_a]["battles"] += 1
    state[model_b]["battles"] += 1
    if sa == 1.0:
        state[model_a]["wins"] += 1
        state[model_b]["losses"] += 1
    elif sa == 0.0:
        state[model_b]["wins"] += 1
        state[model_a]["losses"] += 1
    elif winner != "both_bad":
        state[model_a]["ties"] += 1
        state[model_b]["ties"] += 1

    if winner == "both_bad":
        state[model_a]["both_bad"] = state[model_a].get("both_bad", 0) + 1
        state[model_b]["both_bad"] = state[model_b].get("both_bad", 0) + 1

    save_elo(state, spec)
    return state

def leaderboard_dataframe(state: Optional[Dict[str, dict]] = None, spec: ArenaSpec = MAIN_ARENA) -> pd.DataFrame:
    if state is None:
        state = load_elo(spec)
    rows = []
    for mid in spec.model_ids:
        info = state.get(mid, {"rating": INIT_RATING, "wins": 0, "losses": 0, "battles": 0, "ties": 0})
        p = MODEL_PARAMS.get(mid)
        per_m = round(float(info["rating"]) * 1e6 / p, 1) if p else 0.0
        rows.append({
            "Model": spec.display.get(mid, mid),
            "Model ID": mid,
            "ELO": round(float(info["rating"]), 1),
            "ELO/M params": per_m,
            "Battles": int(info["battles"]),
            "Wins": int(info["wins"]),
            "Losses": int(info["losses"]),
            "Ties": int(info.get("ties", 0)),
            "Both Bad": int(info.get("both_bad", 0)),
        })
    df = pd.DataFrame(rows)
    df = df.sort_values(by="ELO", ascending=False).reset_index(drop=True)
    df.insert(0, "Rank", range(1, len(df) + 1))
    return df


def elo_vs_params_fig(spec: ArenaSpec = MAIN_ARENA) -> "go.Figure":
    """Scatter of ELO rating vs parameter count (log x) for an arena's models."""
    state = load_elo(spec)
    labels, xs, ys, texts, sizes = [], [], [], [], []
    for mid in spec.model_ids:
        p = MODEL_PARAMS.get(mid)
        if not p:
            continue
        info = state.get(mid, {"rating": INIT_RATING, "battles": 0})
        rating = float(info["rating"])
        battles = int(info.get("battles", 0))
        label = spec.display.get(mid, mid)
        labels.append(label)
        xs.append(p)
        ys.append(rating)
        texts.append(f"{label}<br>Params: {p/1e6:.1f}M<br>ELO: {rating:.1f}<br>Battles: {battles}")
        sizes.append(max(8, min(30, 8 + battles * 1.5)))
    fig = go.Figure()
    fig.add_trace(go.Scatter(
        x=xs, y=ys, mode="markers+text",
        text=labels, textposition="top center",
        textfont=dict(size=10, color="#cccccc"),
        marker=dict(size=sizes, color=ys, colorscale="Viridis",
                    line=dict(width=1, color="#ffffff")),
        hovertext=texts, hoverinfo="text",
    ))
    fig.update_layout(
        paper_bgcolor="rgba(0,0,0,0)", plot_bgcolor="rgba(0,0,0,0)",
        font=dict(color="#e0e0e0"),
        xaxis=dict(title="Parameters (log scale)", type="log",
                   gridcolor="rgba(255,255,255,0.08)"),
        yaxis=dict(title="ELO rating", gridcolor="rgba(255,255,255,0.08)"),
        margin=dict(l=50, r=20, t=20, b=50),
        showlegend=False,
    )
    return fig

# ---------------------------------------------------------------------------
# Chat logging to data/chats.jsonl
# ---------------------------------------------------------------------------
def log_battle(spec: ArenaSpec, prompt: str, model_a: str, model_b: str, response_a: str, response_b: str, chosen: str, winner_model: str):
    """
    Append one battle record to data/chats.jsonl.
    Fields: prompt, response_a, response_b, model_a, model_b, chosen (A/B/tie/both_bad), winner_model, timestamp
    Spec: keeps user's message, two responses, each model's names, and what response user chose.
    """
    try:
        get_data_dir().mkdir(parents=True, exist_ok=True)
        record = {
            "timestamp": __import__("datetime").datetime.now(__import__("datetime").timezone.utc).isoformat(),
            "prompt": prompt,
            "model_a": model_a,
            "model_b": model_b,
            "response_a": response_a,
            "response_b": response_b,
            "chosen": chosen,  # "A" / "B" / "tie" / "both_bad"
            "winner_model": winner_model,
            "chosen_response": response_a if chosen == "A" else response_b if chosen == "B" else "",
        }
        with open(get_chat_file(spec.chat_name), "a", encoding="utf-8") as f:
            f.write(json.dumps(record, ensure_ascii=False) + "\n")
    except Exception as e:
        logger.error(f"Failed to log battle: {e}")

# ---------------------------------------------------------------------------
# Model loading (CPU)
# ---------------------------------------------------------------------------
models: Dict[str, object] = {}
tokenizers: Dict[str, object] = {}
model_load_errors: Dict[str, str] = {}

# Diffusion manual instance (if loaded)
diffusion_model: Optional[MetaDiffusionLM] = None
diffusion_tokenizer = None

HF_DIFFUSION_REPO = "CodeSoft/MetaDiffusion-150M-ChatBase"

def load_diffusion_manual():
    """Load MetaDiffusion from HuggingFace (only) using inline architecture."""
    global diffusion_model, diffusion_tokenizer
    if diffusion_model is not None:
        # Re-register in global dicts if cleared (e.g., after tests)
        if "CodeSoft/MetaDiffusion-150M-ChatBase" not in models:
            models["CodeSoft/MetaDiffusion-150M-ChatBase"] = diffusion_model  # type: ignore
        if diffusion_tokenizer is not None and "CodeSoft/MetaDiffusion-150M-ChatBase" not in tokenizers:
            tokenizers["CodeSoft/MetaDiffusion-150M-ChatBase"] = diffusion_tokenizer  # type: ignore
        return diffusion_model, diffusion_tokenizer
    try:
        from huggingface_hub import snapshot_download
        repo_id = HF_DIFFUSION_REPO
        local_dir = Path(snapshot_download(repo_id))
        cfg_path = local_dir / "config.json"
        tok_path = local_dir
        model_path = local_dir / "model.safetensors"
        if not cfg_path.exists() or not model_path.exists():
            logger.warning(f"Diffusion files not found in HF snapshot {local_dir}")
            return None, None
        with open(cfg_path, "r") as f:
            cfg_dict = json.load(f)
        valid = {k: v for k, v in cfg_dict.items() if k in MetaDiffusionConfig.__dataclass_fields__}
        cfg = MetaDiffusionConfig(**valid)
        cfg.tie_word_embeddings = False
        mdl = MetaDiffusionLM(cfg).to(DEVICE)
        from safetensors.torch import load_file
        state = load_file(str(model_path), device="cpu")
        state = {k[len("model."):] if k.startswith("model.") else k: v for k, v in state.items()}
        missing, unexpected = mdl.load_state_dict(state, strict=False)
        if missing or unexpected:
            logger.info(f"  Diffusion load: missing={missing[:3]} unexpected={unexpected[:3]}")
        mdl.to(DEVICE)
        mdl.eval()
        logger.info(f"  Loaded {sum(p.numel() for p in mdl.parameters())/1e6:.1f}M params, vocab={cfg.mask_vocab_size}")
        tok = AutoTokenizer.from_pretrained(str(tok_path), trust_remote_code=True)
        tok = _ensure_diff_chat_tokens(tok)
        if tok.pad_token is None:
            tok.pad_token = tok.eos_token
        diffusion_model = mdl
        diffusion_tokenizer = tok
        logger.info(f"[+] Loaded MetaDiffusion manual from HF {repo_id} (vocab {len(tok)})")
        models["CodeSoft/MetaDiffusion-150M-ChatBase"] = mdl  # type: ignore
        tokenizers["CodeSoft/MetaDiffusion-150M-ChatBase"] = tok  # type: ignore
        return mdl, tok
    except Exception as e:
        logger.warning(f"Manual diffusion load failed: {e}\n{traceback.format_exc()}")
        return None, None

LOCAL_PATHS: Dict[str, str] = {}

def load_models():
    global models, tokenizers, model_load_errors
    # If already populated (including diffusion manual), return
    # But we want to ensure all 5 attempted
    if models and len(models) >= len(MODEL_IDS) + len(BASE_MODEL_IDS) - 1:
        # Already loaded, but ensure diffusion tried
        if "CodeSoft/MetaDiffusion-150M-ChatBase" not in models:
            load_diffusion_manual()
        return models, tokenizers

    logger.info(f"Loading {len(MODEL_IDS) + len(BASE_MODEL_IDS)} models on {DEVICE} ...")
    # Try diffusion manual first (bypass HF Auto which fails on unknown type)
    if "CodeSoft/MetaDiffusion-150M-ChatBase" not in models:
        load_diffusion_manual()

    for mid in MODEL_IDS + BASE_MODEL_IDS:
        if mid in models:
            continue  # already loaded (diffusion)
        load_id = LOCAL_PATHS.get(mid, mid) if os.path.exists(LOCAL_PATHS.get(mid, "")) else mid
        candidates = [load_id]
        if mid in FALLBACK_IDS:
            candidates.append(FALLBACK_IDS[mid])
        success = False
        last_err = None
        for cand in candidates:
            try:
                logger.info(f"[*] Loading {mid} (candidate {cand})...")
                tok = AutoTokenizer.from_pretrained(cand, trust_remote_code=True)
                if tok.pad_token is None:
                    tok.pad_token = tok.eos_token
                mdl = AutoModelForCausalLM.from_pretrained(
                    cand,
                    trust_remote_code=True,
                    torch_dtype=torch.float32,
                    low_cpu_mem_usage=True,
                )
                mdl.to(DEVICE)
                mdl.eval()
                tokenizers[mid] = tok
                models[mid] = mdl
                logger.info(f"[+] Loaded {mid} from {cand} (tok vocab {len(tok)})")
                success = True
                break
            except Exception as e:
                last_err = f"{e}\n{traceback.format_exc()}"
                logger.warning(f"Failed to load {mid} from {cand}: {e}")
                continue
        if not success:
            err_msg = f"Failed candidates {candidates}: {last_err}"
            model_load_errors[mid] = err_msg
            logger.warning(f"[!] {mid} failed to load — generation will error. Error: {err_msg[:600]}")

    logger.info(f"Model loading complete. Loaded: {list(models.keys())} | Failed: {list(model_load_errors.keys())}")
    return models, tokenizers

def ensure_models_loaded():
    # Load if not already attempted
    if not models and not model_load_errors:
        load_models()
    elif "CodeSoft/MetaDiffusion-150M-ChatBase" not in models and not model_load_errors.get("CodeSoft/MetaDiffusion-150M-ChatBase"):
        # Try diffusion again if not yet loaded
        load_diffusion_manual()

# ---------------------------------------------------------------------------
# Prompt formatting & generation
# ---------------------------------------------------------------------------
def build_inputs(tokenizer, model_id: str, prompt: str):
    ctx = MODEL_CONTEXT.get(model_id, 2048)
    gen_budget = GEN_DEFAULTS.get(model_id, {}).get("max_new_tokens", 128)
    max_prompt_tokens = max(32, ctx - gen_budget - 16)
    try:
        if hasattr(tokenizer, "chat_template") and tokenizer.chat_template is not None:
            messages = [{"role": "user", "content": prompt}]
            inputs = tokenizer.apply_chat_template(
                messages, add_generation_prompt=True, return_tensors="pt", truncation=True, max_length=max_prompt_tokens
            )
            if isinstance(inputs, torch.Tensor):
                inputs = {"input_ids": inputs}
            for k in list(inputs.keys()):
                if isinstance(inputs[k], torch.Tensor):
                    inputs[k] = inputs[k].to(DEVICE)
            return inputs
        elif hasattr(tokenizer, "apply_chat_template"):
            try:
                messages = [{"role": "user", "content": prompt}]
                inputs = tokenizer.apply_chat_template(
                    messages, add_generation_prompt=True, return_tensors="pt", truncation=True, max_length=max_prompt_tokens
                )
                if isinstance(inputs, torch.Tensor):
                    inputs = {"input_ids": inputs}
                for k in list(inputs.keys()):
                    if isinstance(inputs[k], torch.Tensor):
                        inputs[k] = inputs[k].to(DEVICE)
                return inputs
            except Exception:
                pass
    except Exception as e:
        logger.debug(f"Chat template failed for {model_id}: {e}")
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=max_prompt_tokens)
    for k in list(inputs.keys()):
        if isinstance(inputs[k], torch.Tensor):
            inputs[k] = inputs[k].to(DEVICE)
    return inputs

def is_diffusion_model(model_id: str) -> bool:
    return "metadiffusion" in model_id.lower()

def generate_for_model(model_id: str, prompt: str, max_new_tokens: int = 0) -> str:
    ensure_models_loaded()
    if model_id not in models or model_id not in tokenizers:
        short = MODEL_DISPLAY.get(model_id, model_id)
        err = model_load_errors.get(model_id, "model not loaded")
        err_short = str(err).splitlines()[0][:800] if err else "model not loaded"
        return f"[Error: {model_id} not loaded: {err_short}]"
    tokenizer = tokenizers[model_id]
    model = models[model_id]
    cfg = GEN_DEFAULTS.get(model_id, {})
    max_new = cfg.get("max_new_tokens", 128)
    if max_new_tokens and int(max_new_tokens) > 0:
        max_new = max(16, min(int(max_new_tokens), 512))
    try:
        if is_diffusion_model(model_id):
            return generate_diffusion(model, tokenizer, prompt, cfg)  # type: ignore
        inputs = build_inputs(tokenizer, model_id, prompt)
        input_len = inputs["input_ids"].shape[1]
        gen_kwargs = {
            "max_new_tokens": max_new,
            "do_sample": cfg.get("do_sample", True),
            "temperature": cfg.get("temperature", 0.7),
            "top_p": cfg.get("top_p", 0.9),
            "repetition_penalty": cfg.get("repetition_penalty", 1.1),
            "pad_token_id": tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id,
            "eos_token_id": tokenizer.eos_token_id,
            "use_cache": False,
        }
        if "top_k" in cfg:
            gen_kwargs["top_k"] = cfg["top_k"]
        if "no_repeat_ngram_size" in cfg:
            gen_kwargs["no_repeat_ngram_size"] = cfg["no_repeat_ngram_size"]
        ctx = MODEL_CONTEXT.get(model_id, 2048)
        if input_len + max_new > ctx:
            gen_kwargs["max_new_tokens"] = max(16, ctx - input_len - 4)
        with torch.inference_mode():
            outputs = model.generate(**inputs, **gen_kwargs)  # type: ignore
        new_tokens = outputs[0, input_len:]
        text = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
        if not text:
            text = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
            prompt_text = tokenizer.decode(inputs["input_ids"][0], skip_special_tokens=True).strip()
            if text.startswith(prompt_text):
                text = text[len(prompt_text):].strip()
        return text if text else "[Empty response]"
    except Exception as e:
        logger.error(f"Generation failed for {model_id}: {e}\n{traceback.format_exc()}")
        return f"[Error generating from {MODEL_DISPLAY.get(model_id, model_id)}: {str(e)[:200]}]"

def generate_diffusion(model, tokenizer, prompt: str, cfg: dict) -> str:
    try:
        tokenizer = _ensure_diff_chat_tokens(tokenizer)
        messages = [{"role": "user", "content": prompt}]
        prompt_str = _format_diff_messages(messages) + f"\n{DIFF_IM_START}assistant\n"
        prompt_ids = torch.tensor([tokenizer.encode(prompt_str, add_special_tokens=False)], device=DEVICE)
        gen_len = int(cfg.get("max_new_tokens", 96))
        num_steps = int(cfg.get("num_steps", 128))
        temperature = float(cfg.get("temperature", 0.7))
        repetition_penalty = float(cfg.get("repetition_penalty", 1.5))
        max_ctx = MODEL_CONTEXT.get("CodeSoft/MetaDiffusion-150M-ChatBase", 5120)
        if prompt_ids.shape[1] + gen_len > max_ctx:
            gen_len = max(16, max_ctx - prompt_ids.shape[1] - 4)
        if gen_len > 256:
            gen_len = 256

        for attempt in range(3):
            cur_temp = temperature * (1 + 0.15 * attempt)
            x = _diff_generate_response(
                model, tokenizer, prompt_ids, gen_len, num_steps, cur_temp, repetition_penalty, DEVICE, stop_on_end=True
            )
            response_tokens = x[0, prompt_ids.shape[1]:].tolist()
            response_tokens = _diff_cut_response(response_tokens, tokenizer)
            text = tokenizer.decode(response_tokens, skip_special_tokens=True).strip()
            if text:
                return text
        return "(empty response)"
    except Exception as e:
        logger.warning(f"Diffusion chat failed: {e}\n{traceback.format_exc()}")
        return f"[Diffusion error] {str(e)[:200]}"

# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------
CSS = """
.gradio-container {max-width: 1450px !important; width: 95% !important;}
.vote-btn {font-weight: 700 !important;}
/* Leaderboard: prevent ELO wrapping, give it fixed width */
#leaderboard, #base_leaderboard { overflow-x: auto; }
#leaderboard table, #base_leaderboard table { table-layout: auto; width: 100%; }
#leaderboard th:nth-child(4), #leaderboard td:nth-child(4),
#base_leaderboard th:nth-child(4), #base_leaderboard td:nth-child(4) {
    min-width: 95px;
    width: 95px;
    white-space: nowrap;
    text-align: center;
    font-variant-numeric: tabular-nums;
}
#leaderboard th:nth-child(1), #leaderboard td:nth-child(1),
#base_leaderboard th:nth-child(1), #base_leaderboard td:nth-child(1) { min-width: 55px; width: 55px; text-align: center; }
#leaderboard td, #base_leaderboard td { white-space: nowrap; overflow: hidden; text-overflow: ellipsis; }
"""

def pick_random_pair(spec: ArenaSpec, exclude_pair: Optional[Tuple[str, str]] = None) -> Tuple[str, str]:
    state = load_elo(spec)
    models_list = spec.model_ids[:]
    weights = []
    C = 5
    K = 100
    for m in models_list:
        games = state.get(m, {}).get("battles", 0)
        w = K / (games + C)
        weights.append(w)
    a = random.choices(models_list, weights=weights, k=1)[0]
    remaining = [m for m in models_list if m != a]
    remaining_weights = [w for m, w in zip(models_list, weights) if m != a]
    b = random.choices(remaining, weights=remaining_weights, k=1)[0]
    if exclude_pair and set((a, b)) == set(exclude_pair):
        a, b = random.sample(spec.model_ids, 2)
    return a, b

def create_demo() -> gr.Blocks:
    with gr.Blocks(title="SLM Arena") as demo:
        gr.Markdown(
            """
            # ⚔️ SLM Arena
            """
        )

        def make_arena_tab(spec: ArenaSpec) -> dict:
            with gr.Tab(spec.arena_title, id=0 if spec.key == "main" else 2):
                prompt = gr.Textbox(
                    label="Your prompt",
                    placeholder="Ask anything... e.g. 'Explain quantum computing in simple terms' or 'Write a haiku about rain'",
                    lines=3,
                    max_length=256,
                )
                max_new = gr.Slider(
                    minimum=64,
                    maximum=512,
                    step=64,
                    value=128,
                    label="Response length (tokens)",
                )
                with gr.Row():
                    submit_btn = gr.Button("⚔️ Battle", variant="primary", scale=1)
                    clear_btn = gr.Button("Clear", variant="secondary", scale=1)
                with gr.Row():
                    with gr.Column():
                        response_a = gr.Textbox(
                            label="Model A", lines=10, max_lines=14, interactive=False,
                            placeholder="Response A will appear here..."
                        )
                        reveal_a = gr.Markdown(visible=False)
                    with gr.Column():
                        response_b = gr.Textbox(
                            label="Model B", lines=10, max_lines=14, interactive=False,
                            placeholder="Response B will appear here..."
                        )
                        reveal_b = gr.Markdown(visible=False)

                with gr.Row():
                    vote_a = gr.Button("👈 Vote for A", variant="secondary", interactive=False, elem_classes=["vote-btn"])
                    vote_tie = gr.Button("🤝 Tie", variant="secondary", interactive=False, elem_classes=["vote-btn"])
                    vote_both_bad = gr.Button("👎 Both Bad", variant="secondary", interactive=False, elem_classes=["vote-btn"])
                    vote_b = gr.Button("Vote for B 👉", variant="secondary", interactive=False, elem_classes=["vote-btn"])

                status = gr.Markdown(visible=False)
                new_round_btn = gr.Button("🔄 New Round", visible=False, variant="secondary")

                model_a_state = gr.State("")
                model_b_state = gr.State("")
                voted_state = gr.State(False)
                prompt_state = gr.State("")

                return {
                    "prompt": prompt,
                    "max_new": max_new,
                    "submit_btn": submit_btn,
                    "clear_btn": clear_btn,
                    "response_a": response_a,
                    "response_b": response_b,
                    "reveal_a": reveal_a,
                    "reveal_b": reveal_b,
                    "vote_a": vote_a,
                    "vote_tie": vote_tie,
                    "vote_both_bad": vote_both_bad,
                    "vote_b": vote_b,
                    "status": status,
                    "new_round_btn": new_round_btn,
                    "model_a_state": model_a_state,
                    "model_b_state": model_b_state,
                    "voted_state": voted_state,
                    "prompt_state": prompt_state,
                    "last_pair": gr.State(None),
                }

        def make_leaderboard_tab(spec: ArenaSpec, elem_id: str, tab_id: int) -> dict:
            with gr.Tab(spec.lb_title, id=tab_id) as tab:
                gr.Markdown("### 🏆 ELO Leaderboard")
                leaderboard = gr.Dataframe(
                    value=leaderboard_dataframe(load_elo(spec), spec),
                    headers=["Rank", "Model", "Model ID", "ELO", "ELO/M params", "Battles", "Wins", "Losses", "Ties", "Both Bad"],
                    datatype=["number", "str", "str", "number", "number", "number", "number", "number", "number", "number"],
                    interactive=False,
                    wrap=False,
                    column_widths=["5%", "15%", "23%", "11%", "11%", "7%", "7%", "7%", "7%", "7%"],
                    elem_id=elem_id,
                )
                gr.Markdown("### 📊 ELO vs Parameters")
                gr.Markdown("X axis is log-scaled; bubble size scales with battle count.")
                elo_params_plot = gr.Plot(value=elo_vs_params_fig(spec), show_label=False)
                refresh_btn = gr.Button("🔄 Refresh", variant="secondary")
                return {"tab": tab, "leaderboard": leaderboard, "elo_params_plot": elo_params_plot, "refresh_btn": refresh_btn}

        with gr.Tabs():
            ui = {}
            lb = {}
            for spec, lb_elem, lb_tab in ((MAIN_ARENA, "leaderboard", 1), (BASE_ARENA, "base_leaderboard", 3)):
                ui[spec.key] = make_arena_tab(spec)
                lb[spec.key] = make_leaderboard_tab(spec, lb_elem, lb_tab)

        # -------------------------------------------------------------------
        # Event handlers
        # -------------------------------------------------------------------
        moderation_model = None
        moderation_tokenizer = None

        def load_moderation():
            nonlocal moderation_model, moderation_tokenizer
            if moderation_model is not None:
                return
            from transformers import BertTokenizer, BertForSequenceClassification
            moderation_tokenizer = BertTokenizer.from_pretrained(MODERATION_MODEL_ID)
            moderation_model = BertForSequenceClassification.from_pretrained(MODERATION_MODEL_ID, num_labels=18)
            moderation_model.to(DEVICE)
            moderation_model.eval()

        def moderate_prompt_impl(prompt: str):
            try:
                load_moderation()
                encoding = moderation_tokenizer(
                    prompt,
                    add_special_tokens=True,
                    max_length=128,
                    padding="max_length",
                    truncation=True,
                    return_attention_mask=True,
                    return_tensors="pt",
                )
                with torch.no_grad():
                    outputs = moderation_model(
                        encoding["input_ids"].to(moderation_model.device),
                        attention_mask=encoding["attention_mask"].to(moderation_model.device),
                    )
                scores = torch.sigmoid(outputs.logits)[0]
                flagged = [MODERATION_CATEGORIES[i] for i in range(len(MODERATION_CATEGORIES)) if scores[i].item() >= MODERATION_THRESHOLD]
                if not flagged:
                    return None
                return ", ".join(flagged)
            except Exception as e:
                logger.error(f"Moderation check failed, allowing prompt: {e}\n{traceback.format_exc()}")
                return None

        @spaces.GPU(duration=120)
        def battle_gpu(spec: ArenaSpec, user_prompt: str, last_pair_val, max_new_tokens: int = 0):
            reason = moderate_prompt_impl(user_prompt)
            if reason is not None:
                return reason, None, None, None, None
            a, b = pick_random_pair(spec, exclude_pair=last_pair_val)
            if random.random() < 0.5:
                a, b = b, a
            ensure_models_loaded()
            resp_a = generate_for_model(a, user_prompt, max_new_tokens)
            resp_b = generate_for_model(b, user_prompt, max_new_tokens)
            return None, resp_a, resp_b, a, b

        def on_submit(spec: ArenaSpec, user_prompt: str, last_pair_val, max_new_tokens: int = 128):
            user_prompt = (user_prompt or "").strip()
            if not user_prompt:
                return (
                    gr.update(value="", placeholder="Please enter a prompt first!"),
                    gr.update(value=""),
                    gr.update(visible=False),
                    gr.update(visible=False),
                    gr.update(visible=False, value=""),
                    gr.update(interactive=False),
                    gr.update(interactive=False),
                    gr.update(interactive=False),
                    gr.update(interactive=False),
                    gr.update(visible=False),
                    "", "", False, user_prompt, last_pair_val,
                    leaderboard_dataframe(load_elo(spec), spec)
                )
            if len(user_prompt) > 256:
                user_prompt = user_prompt[:256]
            reason, resp_a, resp_b, a, b = battle_gpu(spec, user_prompt, last_pair_val, int(max_new_tokens or 128))
            if reason is not None:
                flag_text = f"This prompt was flagged for: {reason}"
                return (
                    gr.update(value=flag_text),
                    gr.update(value=flag_text),
                    gr.update(visible=False),
                    gr.update(visible=False),
                    gr.update(visible=False, value=""),
                    gr.update(interactive=False),
                    gr.update(interactive=False),
                    gr.update(interactive=False),
                    gr.update(interactive=False),
                    gr.update(visible=False),
                    "", "", False, user_prompt, last_pair_val,
                    leaderboard_dataframe(load_elo(spec), spec)
                )
            if not resp_a.strip():
                resp_a = "[No output... model returned empty]"
            if not resp_b.strip():
                resp_b = "[No output... model returned empty]"
            return (
                gr.update(value=resp_a),
                gr.update(value=resp_b),
                gr.update(visible=False),
                gr.update(visible=False),
                gr.update(visible=False, value=""),
                gr.update(interactive=True),
                gr.update(interactive=True),
                gr.update(interactive=True),
                gr.update(interactive=True),
                gr.update(visible=False),
                a, b, False, user_prompt, (a, b),
                leaderboard_dataframe(load_elo(spec), spec)
            )

        def on_vote(spec: ArenaSpec, choice: str, model_a: str, model_b: str, resp_a: str, resp_b: str, user_prompt: str, voted: bool):
            if voted or not model_a or not model_b:
                return (
                    gr.update(visible=False),
                    gr.update(visible=False),
                    gr.update(visible=False, value=""),
                    gr.update(interactive=False),
                    gr.update(interactive=False),
                    gr.update(interactive=False),
                    gr.update(interactive=False),
                    gr.update(visible=False),
                    voted,
                    leaderboard_dataframe(load_elo(spec), spec)
                )
            if choice == "A":
                winner = model_a
                win_label = "A"
                chosen = "A"
            elif choice == "B":
                winner = model_b
                win_label = "B"
                chosen = "B"
            elif choice == "Tie":
                winner = None
                win_label = "Tie"
                chosen = "tie"
            elif choice == "Both Bad":
                winner = "both_bad"
                win_label = "Both Bad"
                chosen = "both_bad"
            else:
                winner = model_b
                win_label = "B"
                chosen = "B"
            state = load_elo(spec)
            ra_before = state[model_a]["rating"]
            rb_before = state[model_b]["rating"]
            update_elo(state, model_a, model_b, winner, spec)
            ra_after = state[model_a]["rating"]
            rb_after = state[model_b]["rating"]
            delta_a = ra_after - ra_before
            delta_b = rb_after - rb_before
            reveal_a_text = f"**Model A:** `{model_a}` ({spec.display.get(model_a, model_a)}) — ELO {ra_after:.1f} ({delta_a:+.1f})"
            reveal_b_text = f"**Model B:** `{model_b}` ({spec.display.get(model_b, model_b)}) — ELO {rb_after:.1f} ({delta_b:+.1f})"
            if choice == "Tie":
                status_text = (
                    f"You voted **Tie**: no winner\n\n"
                    f"**ELO update:** {spec.display.get(model_a, model_a)} {ra_before:.1f} → {ra_after:.1f} ({delta_a:+.1f}) | "
                    f"{spec.display.get(model_b, model_b)} {rb_before:.1f} → {rb_after:.1f} ({delta_b:+.1f})"
                )
            elif choice == "Both Bad":
                status_text = (
                    f"You voted **Both Bad**: no winner\n\n"
                    f"**ELO update:** {spec.display.get(model_a, model_a)} {ra_before:.1f} → {ra_after:.1f} ({delta_a:+.1f}) | "
                    f"{spec.display.get(model_b, model_b)} {rb_before:.1f} → {rb_after:.1f} ({delta_b:+.1f})"
                )
            else:
                status_text = (
                    f"You voted **{win_label}**: the winner is `{winner}`\n\n"
                    f"**ELO update:** {spec.display.get(model_a, model_a)} {ra_before:.1f} → {ra_after:.1f} ({delta_a:+.1f}) | "
                    f"{spec.display.get(model_b, model_b)} {rb_before:.1f} → {rb_after:.1f} ({delta_b:+.1f})"
                )
            # Log chat to data/chats.jsonl
            log_battle(spec, user_prompt, model_a, model_b, resp_a, resp_b, chosen, winner)
            df = leaderboard_dataframe(state, spec)
            return (
                gr.update(value=reveal_a_text, visible=True),
                gr.update(value=reveal_b_text, visible=True),
                gr.update(value=status_text, visible=True),
                gr.update(interactive=False),
                gr.update(interactive=False),
                gr.update(interactive=False),
                gr.update(interactive=False),
                gr.update(visible=True),
                True,
                df
            )

        def on_new_round():
            return (
                gr.update(value=""),
                gr.update(value=""),
                gr.update(value="", visible=False),
                gr.update(value="", visible=False),
                gr.update(value="", visible=False),
                gr.update(interactive=False),
                gr.update(interactive=False),
                gr.update(interactive=False),
                gr.update(interactive=False),
                gr.update(visible=False),
                "", "", False, ""
            )

        def on_clear():
            return (
                gr.update(value=""),
                gr.update(value=""),
                gr.update(value=""),
                gr.update(value="", visible=False),
                gr.update(value="", visible=False),
                gr.update(value="", visible=False),
                gr.update(interactive=False),
                gr.update(interactive=False),
                gr.update(interactive=False),
                gr.update(interactive=False),
                gr.update(visible=False),
                "", "", False, ""
            )

        def on_refresh(spec: ArenaSpec):
            return leaderboard_dataframe(load_elo(spec), spec), elo_vs_params_fig(spec)

        for spec in (MAIN_ARENA, BASE_ARENA):
            u, b = ui[spec.key], lb[spec.key]
            round_outputs = [u["response_a"], u["response_b"], u["reveal_a"], u["reveal_b"], u["status"], u["vote_a"], u["vote_tie"], u["vote_both_bad"], u["vote_b"], u["new_round_btn"], u["model_a_state"], u["model_b_state"], u["voted_state"], u["prompt_state"]]
            submit_outputs = round_outputs + [u["last_pair"], b["leaderboard"]]
            vote_outputs = [u["reveal_a"], u["reveal_b"], u["status"], u["vote_a"], u["vote_tie"], u["vote_both_bad"], u["vote_b"], u["new_round_btn"], u["voted_state"], b["leaderboard"]]

            u["submit_btn"].click(
                fn=lambda p, lp, mn, s=spec: on_submit(s, p, lp, mn),
                inputs=[u["prompt"], u["last_pair"], u["max_new"]],
                outputs=submit_outputs,
            )
            u["prompt"].submit(
                fn=lambda p, lp, mn, s=spec: on_submit(s, p, lp, mn),
                inputs=[u["prompt"], u["last_pair"], u["max_new"]],
                outputs=submit_outputs,
            )
            for btn, choice in ((u["vote_a"], "A"), (u["vote_tie"], "Tie"), (u["vote_both_bad"], "Both Bad"), (u["vote_b"], "B")):
                btn.click(
                    fn=lambda ma, mb, ra, rb, pr, vd, c=choice, s=spec: on_vote(s, c, ma, mb, ra, rb, pr, vd),
                    inputs=[u["model_a_state"], u["model_b_state"], u["response_a"], u["response_b"], u["prompt_state"], u["voted_state"]],
                    outputs=vote_outputs,
                )
            u["new_round_btn"].click(fn=on_new_round, inputs=[], outputs=round_outputs)
            u["clear_btn"].click(fn=on_clear, inputs=[], outputs=[u["prompt"]] + round_outputs)
            b["refresh_btn"].click(fn=lambda s=spec: on_refresh(s), inputs=[], outputs=[b["leaderboard"], b["elo_params_plot"]])
            b["tab"].select(fn=lambda s=spec: on_refresh(s), inputs=[], outputs=[b["leaderboard"], b["elo_params_plot"]])

        gr.Markdown(
            "<div style=\"text-align:center; font-size:12px; opacity:0.65;\">"
            "All prompts and responses are logged and shared publicly. "
            "Do not enter any harmful or personal information. "
            "Models may output incorrect, harmful, or offensive responses."
            "</div>"
        )

        try:
            demo.load(fn=lambda: on_refresh(MAIN_ARENA), inputs=[], outputs=[lb["main"]["leaderboard"], lb["main"]["elo_params_plot"]], show_progress="hidden")
        except Exception:
            pass

    return demo

# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    print("=" * 60)
    print(f"SLM Arena starting, attempting to load {len(MODEL_IDS) + len(BASE_MODEL_IDS)} models on {DEVICE}...")
    print(f"Models: {MODEL_IDS + BASE_MODEL_IDS}")
    print(f"Data dir: {get_data_dir().resolve()} (bucket /data if mounted)")
    print("=" * 60)
    try:
        load_models()
    except Exception as e:
        logger.error(f"Model loading encountered error: {e}")
    try:
        for spec in (MAIN_ARENA, BASE_ARENA):
            df = leaderboard_dataframe(load_elo(spec), spec)
            print(df.to_string(index=False))
            chat_path = get_chat_file(spec.chat_name)
            print(f"\nChat log: {chat_path.resolve()} (exists={chat_path.exists()})")
            if chat_path.exists():
                with open(chat_path) as f:
                    lines = sum(1 for _ in f)
            else:
                lines = 0
            print(f"Previous battles logged: {lines}")
    except Exception as e:
        logger.warning(f"Leaderboard preview failed: {e}")
    demo = create_demo()
    demo.queue(max_size=20)
    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True, theme=gr.themes.Base(), css=CSS)