app.py

#!/usr/bin/env python3
# app.py — Chat inference for AGILLM2 (HF Spaces friendly)
# - Auto-detects non-interactive env (Spaces) and launches Gradio UI
# - Loads final.pt from repo OpenTransformer/AGILLM2-fast-training
# - Qwen tokenizer + chat template
# - Optional local CLI REPL when run in a terminal
# - Adds a "Raw transcript" tab with "User:" / "Assistant:" lines

from __future__ import annotations
import os, sys, time, math, argparse
from typing import Optional, Tuple, List, Dict, Any

import torch
import torch.nn as nn
import torch.nn.functional as F

# Silence the PyTorch TF32 deprecation nag in Spaces logs (optional)
import warnings
warnings.filterwarnings("ignore", message="Please use the new API settings to control TF32 behavior")

from huggingface_hub import hf_hub_download
from transformers import AutoTokenizer, logging as hf_log
hf_log.set_verbosity_error()

# ───────────────────────── Config ─────────────────────────
MODEL_REPO   = os.getenv("MODEL_REPO", "OpenTransformer/AGILLM2-fast-training")
CKPT_NAME    = os.getenv("CKPT_NAME",  "final.pt")  # e.g., step04121612.pt
TOKENIZER_ID = os.getenv("TOKENIZER_ID", "Qwen/Qwen3-235B-A22B-Thinking-2507")

DEV = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.backends.cuda.matmul.allow_tf32 = True
try:
    torch.set_float32_matmul_precision("high")
except Exception:
    pass

# ───────────────────────── Tokenizer ─────────────────────────
tok = AutoTokenizer.from_pretrained(TOKENIZER_ID, use_fast=True, trust_remote_code=True)
if tok.pad_token is None:
    tok.add_special_tokens({"pad_token": "[PAD]"})
VOCAB = max(tok.get_vocab().values()) + 1
BLANK = tok.pad_token_id
EOS = tok.eos_token_id if tok.eos_token_id is not None else tok.sep_token_id

# ───────────────────────── AMP helper ─────────────────────────
try:
    from torch.amp import autocast as _ac, GradScaler  # noqa
except Exception:
    from torch.cuda.amp import autocast as _ac, GradScaler  # noqa

def _supports_fp8() -> bool:
    return hasattr(torch, "float8_e4m3fn")

def _auto_amp_dtype(prefer_fp8: bool = False):
    if DEV.type != "cuda":
        return torch.float32
    if prefer_fp8 and _supports_fp8():
        return torch.float8_e4m3fn
    try:
        if torch.cuda.is_bf16_supported():
            return torch.bfloat16
        return torch.float16
    except Exception:
        return torch.float16

def amp(enabled: bool, prefer_fp8: bool = False):
    if not (enabled and DEV.type == "cuda"):
        from contextlib import nullcontext
        return nullcontext()
    return _ac(device_type="cuda", dtype=_auto_amp_dtype(prefer_fp8=prefer_fp8))

# ───────────────────────── ALiBi helpers ─────────────────────────
def _alibi_slopes(n_heads: int):
    import math as _m
    def pow2slopes(n):
        start = 2 ** (-2 ** -(_m.log2(n) - 3))
        ratio = start
        return [start * (ratio ** i) for i in range(n)]
    if _m.log2(n_heads).is_integer():
        vals = pow2slopes(n_heads)
    else:
        closest = 2 ** _m.floor(_m.log2(n_heads))
        vals = pow2slopes(closest); extra = pow2slopes(2 * closest)
        vals += extra[0::2][: n_heads - closest]
    return torch.tensor(vals, device=DEV).view(1, n_heads, 1, 1)

def alibi_bias(n_heads: int, n_tokens: int):
    i = torch.arange(n_tokens, device=DEV).view(1, 1, n_tokens, 1)
    j = torch.arange(n_tokens, device=DEV).view(1, 1, 1, n_tokens)
    dist = (j - i).clamp_min(0)
    slopes = _alibi_slopes(n_heads)
    return -slopes * dist

# ───────────────────────── Model (Encoder + AR head) ─────────────────────────
class LowRankMHA(nn.Module):
    def __init__(self, d: int, h: int, r: int, use_relpos: bool = True):
        super().__init__()
        assert d % h == 0, "d must be divisible by number of heads"
        self.h, self.dk = h, d // h
        self.use_relpos = use_relpos
        self.q = nn.Linear(d, d, bias=False)
        self.k = nn.Linear(d, d, bias=False)
        self.v = nn.Linear(d, d, bias=False)
        self.U = nn.Parameter(torch.randn(self.dk, r))
        nn.init.orthogonal_(self.U)
        self.proj = nn.Linear(h * r, d, bias=False)
        self.drop = nn.Dropout(0.1)

    def _proj(self, x):
        B, N, _ = x.shape
        return (x.view(B, N, self.h, self.dk).transpose(1, 2) @ self.U)

    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None,
                rel_bias_tokens: Optional[int] = None,
                kv_cache: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
                use_cache: bool = False):
        q = self._proj(self.q(x))
        k_new = self._proj(self.k(x))
        v_new = self._proj(self.v(x))
        if kv_cache is None:
            k, v = k_new, v_new
        else:
            k, v = kv_cache
            if use_cache:
                k = torch.cat([k, k_new], dim=2)
                v = torch.cat([v, v_new], dim=2)
        att = (q @ k.transpose(-1, -2)) / math.sqrt(self.dk)
        if q.size(2) == k.size(2):
            if self.use_relpos and rel_bias_tokens is not None:
                att = att + alibi_bias(self.h, rel_bias_tokens)
            if mask is not None:
                att = att + mask
        z = (att.softmax(-1) @ v).transpose(1, 2)
        z = z.reshape(x.size(0), x.size(1), -1)
        out = self.drop(self.proj(z))
        return (out, (k, v)) if use_cache else out

class Block(nn.Module):
    def __init__(self, d: int, h: int, r: int):
        super().__init__()
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.mha = LowRankMHA(d, h, r, use_relpos=True)
        self.ff = nn.Sequential(nn.Linear(d, 4 * d), nn.ReLU(), nn.Linear(4 * d, d))

    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor],
                kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
                use_cache: bool = False):
        n = x.size(1)
        if use_cache:
            y, new_kv = self.mha(self.ln1(x), mask, rel_bias_tokens=n if mask is not None else None, kv_cache=kv, use_cache=True)
            x = x + y
            x = x + self.ff(self.ln2(x))
            return x, new_kv
        else:
            x = x + self.mha(self.ln1(x), mask, rel_bias_tokens=n)
            return x + self.ff(self.ln2(x))

class Encoder(nn.Module):
    def __init__(self, cfg: Dict[str, int]):
        super().__init__()
        d, l, h, r = cfg["d"], cfg["layers"], cfg["heads"], cfg["rank"]
        self.emb = nn.Embedding(VOCAB, d)
        self.blocks = nn.ModuleList([Block(d, h, r) for _ in range(l)])
        self.ln = nn.LayerNorm(d)

    def forward(self, ids: torch.Tensor, mask: Optional[torch.Tensor],
                kv_caches: Optional[List[Optional[Tuple[torch.Tensor, torch.Tensor]]]] = None,
                use_cache: bool = False):
        x = self.emb(ids)
        if not use_cache:
            for blk in self.blocks:
                x = blk(x, mask)
            return self.ln(x)
        new_kvs: List[Tuple[torch.Tensor, torch.Tensor]] = []
        for i, blk in enumerate(self.blocks):
            kv = kv_caches[i] if (kv_caches is not None) else None
            x, kv_out = blk(x, mask, kv, use_cache=True)
            new_kvs.append(kv_out)
        return self.ln(x), new_kvs

class ARHead(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.proj = nn.Linear(d, VOCAB)
    def forward(self, h): return self.proj(h)

# ───────────────────────── Misc ─────────────────────────
def causal_mask(n: int):
    m = torch.full((1, 1, n, n), float("-inf"), device=DEV)
    return torch.triu(m, 1)

def _resolve_cfg_from_ckpt(sd: dict) -> Dict[str, int]:
    if isinstance(sd, dict) and "cfg" in sd and isinstance(sd["cfg"], dict):
        return dict(sd["cfg"])
    core = sd.get("core", {})
    emb_w = core.get("emb.weight")
    if emb_w is None:
        raise RuntimeError("Checkpoint missing core.emb.weight; cannot infer d/l/h/r.")
    d = emb_w.shape[1]
    layer_ids = []
    for k in core.keys():
        if k.startswith("blocks."):
            parts = k.split(".")
            if len(parts) > 2 and parts[1].isdigit():
                layer_ids.append(int(parts[1]))
    layers = (max(layer_ids) + 1) if layer_ids else 0
    U = core.get("blocks.0.mha.U")
    if U is None:
        raise RuntimeError("Checkpoint missing blocks.0.mha.U; cannot infer rank/heads.")
    dk, r = U.shape
    h = d // dk
    return {"d": d, "layers": layers, "heads": h, "rank": r}

def load_joint_from_hub(repo_id: str, filename: str):
    ckpt_path = hf_hub_download(repo_id=repo_id, filename=filename, resume_download=True)
    sd = torch.load(ckpt_path, map_location="cpu")
    cfg = _resolve_cfg_from_ckpt(sd)
    core = Encoder(cfg).to(DEV)
    ar_h = ARHead(cfg["d"]).to(DEV)
    core.load_state_dict(sd["core"])
    if "ar" in sd: ar_h.load_state_dict(sd["ar"])
    return core, ar_h, cfg

# ───────────────────────── Chat helpers ─────────────────────────
def render_chat(messages: List[Dict[str, str]], add_generation_prompt: bool = True) -> str:
    try:
        return tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=add_generation_prompt)
    except Exception:
        # Fallback plain format
        out = []
        for m in messages:
            role = m.get("role", "user")
            content = m.get("content", "")
            out.append(f"{role.capitalize()}: {content}")
        if add_generation_prompt:
            out.append("Assistant:")
        return "\n".join(out)

def render_raw(history: List[Tuple[str, str]] | None, sys_prompt: str) -> str:
    lines = []
    if sys_prompt:
        lines.append(f"System: {sys_prompt}")
    for u, a in (history or []):
        lines.append(f"User: {u}")
        lines.append(f"Assistant: {a}")
    return "\n".join(lines)

def _apply_no_repeat_ngram(logits: torch.Tensor, ids: torch.Tensor, n: int):
    if n <= 0 or ids.size(1) < n - 1: return logits
    prefix = ids[0, -(n - 1):].tolist()
    banned, tokens = [], ids[0].tolist()
    for i in range(len(tokens) - n + 1):
        if tokens[i:i + n - 1] == prefix:
            banned.append(tokens[i + n - 1])
    if banned:
        banned_idx = torch.tensor(banned, device=logits.device, dtype=torch.long)
        logits[..., banned_idx] = float("-inf")
    return logits

def _apply_rep_presence_frequency(logits, ids, last_n, repetition_penalty, presence_penalty, frequency_penalty):
    if ids.numel() == 0: return logits
    hist = ids[0, -last_n:].to(torch.long) if last_n > 0 else ids[0].to(torch.long)
    if hist.numel() == 0: return logits
    uniq, counts = torch.unique(hist, return_counts=True)
    if presence_penalty != 0.0 or frequency_penalty != 0.0:
        adjust = presence_penalty + frequency_penalty * counts.to(logits.dtype)
        logits[..., uniq] = logits[..., uniq] - adjust
    if repetition_penalty and abs(repetition_penalty - 1.0) > 1e-6:
        sel = logits[..., uniq]
        sel = torch.where(sel > 0, sel / repetition_penalty, sel * repetition_penalty)
        logits[..., uniq] = sel
    return logits

def _filter_top_k_top_p_min_p(logits: torch.Tensor, top_k: int, top_p: float, min_p: float, temperature: float):
    logits = logits / max(temperature, 1e-8)
    if logits.dim() == 1: logits = logits.unsqueeze(0)
    probs = logits.softmax(-1)
    V = probs.size(-1)
    if top_k and top_k < V:
        vals, idx = torch.topk(probs, top_k, dim=-1)
        mask = torch.full_like(probs, 0.0); mask.scatter_(1, idx, 1.0); probs = probs * mask
    if top_p < 1.0:
        sorted_probs, sorted_idx = torch.sort(probs, descending=True, dim=-1)
        cumsum = torch.cumsum(sorted_probs, dim=-1)
        keep = cumsum <= top_p; keep[..., 0] = True
        mask = torch.zeros_like(probs); mask.scatter_(1, sorted_idx, keep.to(mask.dtype))
        probs = probs * mask
    if min_p > 0.0:
        probs = torch.where(probs >= min_p, probs, torch.zeros_like(probs))
    sums = probs.sum(-1, keepdim=True); empty = (sums == 0)
    if empty.any():
        fallback_idx = logits.argmax(-1, keepdim=True)
        probs = torch.where(empty, torch.zeros_like(probs), probs)
        probs.scatter_(-1, fallback_idx, torch.where(empty, torch.ones_like(sums), torch.zeros_like(sums)))
    probs = probs / probs.sum(-1, keepdim=True)
    return probs

@torch.no_grad()
def chat_decode(core, ar_h, messages: List[Dict[str, str]], max_new: int = 200, T: float = 0.9,
                greedy: bool = False, top_k: int = 50, top_p: float = 0.9, min_p: float = 0.0,
                repetition_penalty: float = 1.1, presence_penalty: float = 0.3, frequency_penalty: float = 0.2,
                penalty_last_n: int = 128, no_repeat_ngram_size: int = 3,
                use_fp8: bool = False, fp8_fallback: bool = True) -> str:
    prompt = render_chat(messages, add_generation_prompt=True)
    ids = torch.tensor([tok.encode(prompt)], device=DEV)
    prompt_len = ids.size(1)

    with amp(use_fp8 or False, prefer_fp8=(use_fp8 and (_supports_fp8() or fp8_fallback))):
        h_full, kvs = core(ids, causal_mask(ids.size(1)), use_cache=True)
        for _ in range(max_new):
            logits = ar_h(h_full)[:, -1]
            logits = _apply_no_repeat_ngram(logits, ids, no_repeat_ngram_size)
            logits = _apply_rep_presence_frequency(logits, ids, penalty_last_n,
                                                   repetition_penalty, presence_penalty, frequency_penalty)
            if greedy:
                nxt = logits.argmax(-1, keepdim=True)
            else:
                probs = _filter_top_k_top_p_min_p(logits.squeeze(0), top_k, top_p, min_p, T)
                nxt = probs.multinomial(1)
            ids = torch.cat([ids, nxt.unsqueeze(0) if nxt.dim()==1 else nxt], 1)
            x = ids[:, -1:]
            h_full, kvs = core(x, None, kv_caches=kvs, use_cache=True)

    full_ids = ids[0].tolist()
    return tok.decode(full_ids[prompt_len:], skip_special_tokens=True).strip()

# ───────────────────────── UI / CLI ─────────────────────────
def launch_gradio(core, ar_h):
    import gradio as gr
    with gr.Blocks() as demo:
        gr.Markdown("### OpenTransformer / AGILLM2 — Chat")

        with gr.Row():
            temp = gr.Slider(0.1, 1.5, value=0.9, step=0.05, label="Temperature")
            topp = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
            topk = gr.Slider(0, 200, value=50, step=1, label="Top-k")
            mxnt = gr.Slider(16, 1024, value=200, step=8, label="Max new tokens")
            sys_prompt = gr.Textbox(value="You are a helpful, concise assistant.", label="System prompt")

        with gr.Tabs():
            with gr.TabItem("Chat"):
                chatbot = gr.Chatbot(height=520, label="Conversation")
                msg = gr.Textbox(placeholder="Say something useful…", label="Message")
                submit = gr.Button("Send", variant="primary")
            with gr.TabItem("Raw transcript"):
                raw = gr.Textbox(lines=24, label="Raw transcript (User:/Assistant:)", interactive=False)

        clear = gr.Button("Clear", variant="secondary")

        def _chat(history, user_msg, t, p, k, mxt, sys_p):
            # Build messages from history + new user message
            messages = [{"role":"system","content":sys_p}]
            for u,a in history or []:
                messages.append({"role":"user","content":u})
                messages.append({"role":"assistant","content":a})
            messages.append({"role":"user","content":user_msg})

            reply = chat_decode(
                core, ar_h, messages,
                max_new=int(mxt), T=float(t),
                greedy=False, top_k=int(k), top_p=float(p),
                use_fp8=False, fp8_fallback=True
            )
            history = (history or []) + [(user_msg, reply)]
            transcript = render_raw(history, sys_p)
            return history, "", transcript

        # Wire up events: submit via button or enter
        msg.submit(_chat, [chatbot, msg, temp, topp, topk, mxnt, sys_prompt], [chatbot, msg, raw], queue=False)
        submit.click(_chat, [chatbot, msg, temp, topp, topk, mxnt, sys_prompt], [chatbot, msg, raw], queue=False)

        def _clear():
            return [], "", ""
        clear.click(_clear, inputs=None, outputs=[chatbot, msg, raw], queue=False)

    demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))

def run_cli(core, ar_h):
    print("Type to chat. Ctrl+C to exit.")
    history: List[Tuple[str,str]] = []
    while True:
        try:
            user = input("\nYou: ").strip()
            if not user: continue
            messages = [{"role":"system","content":"You are a helpful, concise assistant."}]
            for u,a in history:
                messages.append({"role":"user","content":u})
                messages.append({"role":"assistant","content":a})
            messages.append({"role":"user","content":user})
            t0 = time.time()
            reply = chat_decode(core, ar_h, messages, max_new=200, T=0.9, top_k=50, top_p=0.9)
            dt = time.time()-t0
            print(f"Bot: {reply}\n[{len(tok.encode(reply))} tok in {dt:.2f}s]")
            history.append((user, reply))
            # Also show raw transcript line by line in CLI
            print("\n--- RAW ---")
            print(render_raw(history, "You are a helpful, concise assistant."))
        except KeyboardInterrupt:
            print("\nbye."); break

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--cli", action="store_true", help="Force CLI REPL even if not a TTY")
    ap.add_argument("--gradio", action="store_true", help="Force Gradio UI")
    args = ap.parse_args()

    print(f"[init] downloading checkpoint {CKPT_NAME} from {MODEL_REPO} …", flush=True)
    core, ar_h, cfg = load_joint_from_hub(MODEL_REPO, CKPT_NAME)
    core.eval(); ar_h.eval()
    print(f"[ready] cfg={cfg} device={DEV.type} vocab={VOCAB}", flush=True)

    # Spaces have no interactive stdin. Auto-launch Gradio if not a TTY.
    in_tty = sys.stdin.isatty()
    if args.gradio or (not args.cli and not in_tty):
        launch_gradio(core, ar_h)
    else:
        run_cli(core, ar_h)

if __name__ == "__main__":
    main()