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", "CodeSoft/sorbet-v2-25m", "DALabCommunity/Haidass1.5-143M", ] 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", "CodeSoft/sorbet-v2-25m": "sorbet-v2-25m", "DALabCommunity/Haidass1.5-143M": "Haidass1.5-143M", } 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, "CodeSoft/sorbet-v2-25m": 25.2e6, "DALabCommunity/Haidass1.5-143M": 143.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}, "CodeSoft/sorbet-v2-25m": {"max_new_tokens": 64, "temperature": 0.8, "top_p": 0.95, "repetition_penalty": 1.1, "do_sample": True}, "DALabCommunity/Haidass1.5-143M": {"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, "CodeSoft/sorbet-v2-25m": 4096, "DALabCommunity/Haidass1.5-143M": 4096, } @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 ; 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}
Params: {p/1e6:.1f}M
ELO: {rating:.1f}
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( "
" "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." "
" ) 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)