Delete app.py

app.py

@@ -1,437 +0,0 @@
-"""Cosmos-T2-80M-Test — Gradio Chat Demo
-
-Standalone inference app generated by the Cosmos-T2 universal training notebook.
-It matches the notebook architecture: RoPE, RMSNorm, SwiGLU, GQA, and a configurable Engram memory path.
-"""
-
-import contextlib
-import math
-from pathlib import Path
-
-import gradio as gr
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from huggingface_hub import hf_hub_download
-from transformers import AutoTokenizer
-
-MODEL_REPO_ID = "wop/Cosmos-T2-80M-Test"
-CHECKPOINT_NAME = "Cosmos-T2-80M-Test.pt"
-TOKENIZER_NAME = "Qwen/Qwen2.5-0.5B"
-MODEL_NAME = "Cosmos-T2-80M-Test"
-DEFAULT_SYSTEM_PROMPT = "Enable thinking features: INTUITION, COLD START, HOT START"
-MAX_CTX_HARD = 1028
-MAX_NEW_HARD = 1028
-
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
-
-class RMSNorm(nn.Module):
-    def __init__(self, dim, eps=1e-6):
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(dim))
-        self.eps = eps
-    def forward(self, x):
-        rms = x.pow(2).mean(dim=-1, keepdim=True)
-        x = x * torch.rsqrt(rms + self.eps)
-        return x * self.weight
-
-def rotate_half(x):
-    x1 = x[..., ::2]
-    x2 = x[..., 1::2]
-    return torch.stack((-x2, x1), dim=-1).flatten(-2)
-
-def apply_rope(q, k, cos, sin):
-    return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)
-
-class GQAAttention(nn.Module):
-    def __init__(self, d_model, n_heads, n_kv_heads, rope_base=10000, dropout=0.0):
-        super().__init__()
-        assert d_model % n_heads == 0
-        assert n_heads % n_kv_heads == 0
-        self.n_heads = n_heads
-        self.n_kv_heads = n_kv_heads
-        self.head_dim = d_model // n_heads
-        self.dropout = dropout
-        self.q_proj = nn.Linear(d_model, n_heads * self.head_dim, bias=False)
-        self.k_proj = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)
-        self.v_proj = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)
-        self.o_proj = nn.Linear(d_model, d_model, bias=False)
-
-    def forward(self, x, rope_cos, rope_sin):
-        batch, seq_len, _ = x.shape
-        q = self.q_proj(x).view(batch, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
-        k = self.k_proj(x).view(batch, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2)
-        v = self.v_proj(x).view(batch, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2)
-        q, k = apply_rope(q, k, rope_cos, rope_sin)
-        if self.n_kv_heads != self.n_heads:
-            repeat = self.n_heads // self.n_kv_heads
-            k = k.repeat_interleave(repeat, dim=1)
-            v = v.repeat_interleave(repeat, dim=1)
-        out = F.scaled_dot_product_attention(
-            q, k, v, is_causal=True, dropout_p=self.dropout if self.training else 0.0
-        )
-        out = out.transpose(1, 2).contiguous().view(batch, seq_len, -1)
-        return self.o_proj(out)
-
-class SwiGLUMLP(nn.Module):
-    def __init__(self, d_model, hidden_dim, dropout=0.0):
-        super().__init__()
-        self.gate = nn.Linear(d_model, hidden_dim, bias=False)
-        self.up = nn.Linear(d_model, hidden_dim, bias=False)
-        self.down = nn.Linear(hidden_dim, d_model, bias=False)
-        self.dropout = nn.Dropout(dropout)
-    def forward(self, x):
-        return self.down(self.dropout(F.silu(self.gate(x)) * self.up(x)))
-
-class EngramMemory(nn.Module):
-    def __init__(self, d_model, bucket_count, memory_dim, order, pad_id=0, dropout=0.0):
-        super().__init__()
-        self.bucket_count = bucket_count
-        self.order = order
-        self.pad_id = pad_id
-        self.bucket = nn.Embedding(bucket_count, memory_dim)
-        self.query = nn.Linear(d_model, memory_dim, bias=False)
-        self.project = nn.Linear(memory_dim, d_model, bias=False)
-        self.gate = nn.Linear(d_model, d_model, bias=True)
-        self.dropout = nn.Dropout(dropout)
-        primes = [1, 1315423911, 2654435761, 97531, 433494437]
-        self.register_buffer("primes", torch.tensor(primes[:order], dtype=torch.long), persistent=False)
-
-    def hash_tokens(self, idx):
-        batch, seq_len = idx.shape
-        pad = torch.full((batch, self.order - 1), self.pad_id, device=idx.device, dtype=idx.dtype)
-        history = torch.cat([pad, idx], dim=1)
-        hashed = torch.zeros((batch, seq_len), device=idx.device, dtype=torch.long)
-        for offset in range(self.order):
-            slice_ = history[:, offset: offset + seq_len].long()
-            hashed = (hashed * 1315423911 + slice_ * self.primes[offset]) % self.bucket_count
-        return hashed
-
-    def forward(self, x, idx):
-        hashed = self.hash_tokens(idx)
-        query = torch.tanh(self.query(x))
-        memory = self.bucket(hashed) * query
-        memory = self.project(memory)
-        gate = torch.sigmoid(self.gate(x))
-        return self.dropout(gate * memory)
-
-class Block(nn.Module):
-    def __init__(self, d_model, n_heads, n_kv_heads, d_ff, rope_base, dropout=0.0, use_engram=False, engram_bucket_count=4096, engram_dim=96, engram_order=3, pad_id=0):
-        super().__init__()
-        self.norm1 = RMSNorm(d_model)
-        self.attn = GQAAttention(d_model, n_heads, n_kv_heads, rope_base=rope_base, dropout=dropout)
-        self.norm2 = RMSNorm(d_model)
-        self.engram = EngramMemory(d_model, engram_bucket_count, engram_dim, engram_order, pad_id=pad_id, dropout=dropout) if use_engram else None
-        self.norm3 = RMSNorm(d_model)
-        self.mlp = SwiGLUMLP(d_model, d_ff, dropout=dropout)
-    def forward(self, x, idx, rope_cos, rope_sin):
-        x = x + self.attn(self.norm1(x), rope_cos, rope_sin)
-        if self.engram is not None:
-            x = x + self.engram(self.norm2(x), idx)
-        return x + self.mlp(self.norm3(x))
-
-class CosmosT2_LLM(nn.Module):
-    def __init__(self, vocab_size, d_model=384, n_layers=12, n_heads=8, n_kv_heads=2, d_ff=1536, max_len=1028, rope_base=10000, dropout=0.05, use_engram=True, engram_every=2, engram_bucket_count=4096, engram_dim=96, engram_order=3, pad_id=0):
-        super().__init__()
-        self.vocab_size = vocab_size
-        self.d_model = d_model
-        self.n_layers = n_layers
-        self.n_heads = n_heads
-        self.n_kv_heads = n_kv_heads
-        self.head_dim = d_model // n_heads
-        self.max_len = max_len
-        self.rope_base = rope_base
-        self.pad_id = pad_id
-        self.tok_emb = nn.Embedding(vocab_size, d_model)
-        self.blocks = nn.ModuleList()
-        for layer_index in range(n_layers):
-            block_uses_engram = use_engram and ((layer_index + 1) % engram_every == 0)
-            self.blocks.append(
-                Block(
-                    d_model=d_model,
-                    n_heads=n_heads,
-                    n_kv_heads=n_kv_heads,
-                    d_ff=d_ff,
-                    rope_base=rope_base,
-                    dropout=dropout,
-                    use_engram=block_uses_engram,
-                    engram_bucket_count=engram_bucket_count,
-                    engram_dim=engram_dim,
-                    engram_order=engram_order,
-                    pad_id=pad_id,
-                )
-            )
-        self.norm_f = RMSNorm(d_model)
-        self.apply(self._init_weights)
-
-    def _init_weights(self, module):
-        if isinstance(module, nn.Linear):
-            nn.init.normal_(module.weight, mean=0.0, std=0.02)
-            if module.bias is not None:
-                nn.init.zeros_(module.bias)
-        elif isinstance(module, nn.Embedding):
-            nn.init.normal_(module.weight, mean=0.0, std=0.02)
-
-    def build_rope(self, seq_len, device, dtype):
-        inv_freq = 1.0 / (self.rope_base ** (torch.arange(0, self.head_dim, 2, device=device, dtype=torch.float32) / self.head_dim))
-        positions = torch.arange(seq_len, device=device, dtype=torch.float32)
-        freqs = torch.outer(positions, inv_freq)
-        cos = freqs.cos().to(dtype)[None, None, :, :]
-        sin = freqs.sin().to(dtype)[None, None, :, :]
-        return cos, sin
-
-    def forward(self, idx, targets=None):
-        if idx.size(1) > self.max_len:
-            idx = idx[:, -self.max_len:]
-        seq_len = idx.size(1)
-        rope_cos, rope_sin = self.build_rope(seq_len, idx.device, self.tok_emb.weight.dtype)
-        x = self.tok_emb(idx)
-        for block in self.blocks:
-            x = block(x, idx, rope_cos, rope_sin)
-        x = self.norm_f(x)
-        logits = F.linear(x, self.tok_emb.weight)
-        loss = None
-        if targets is not None:
-            loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1))
-        return logits, loss
-
-    @torch.no_grad()
-    def generate_step(self, idx, temperature=0.8, top_k=50, max_ctx=None):
-        max_ctx = self.max_len if max_ctx is None else max_ctx
-        idx_cond = idx[:, -max_ctx:]
-        logits, _ = self(idx_cond)
-        logits = logits[:, -1, :]
-        if temperature <= 1e-6:
-            return torch.argmax(logits, dim=-1, keepdim=True)
-        logits = logits / temperature
-        if top_k and top_k > 0:
-            values, _ = torch.topk(logits, min(top_k, logits.size(-1)))
-            cutoff = values[:, [-1]]
-            logits = logits.masked_fill(logits < cutoff, float("-inf"))
-        probs = F.softmax(logits, dim=-1)
-        return torch.multinomial(probs, num_samples=1)
-
-    @torch.no_grad()
-    def generate(self, idx, max_new_tokens=128, temperature=0.8, top_k=50, max_ctx=None):
-        for _ in range(max_new_tokens):
-            nxt = self.generate_step(idx, temperature=temperature, top_k=top_k, max_ctx=max_ctx)
-            idx = torch.cat([idx, nxt], dim=1)
-        return idx
-
-def load_checkpoint(tokenizer):
-    ckpt_path = Path(CHECKPOINT_NAME)
-    if not ckpt_path.exists():
-        print(f"Downloading {CHECKPOINT_NAME} from {MODEL_REPO_ID}")
-        ckpt_path = Path(hf_hub_download(repo_id=MODEL_REPO_ID, filename=CHECKPOINT_NAME))
-    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
-    cfg = ckpt.get("config", {})
-    resolved = {
-        "vocab_size": cfg.get("vocab_size", len(tokenizer)),
-        "d_model": cfg.get("d_model", 384),
-        "n_layers": cfg.get("n_layers", 12),
-        "n_heads": cfg.get("n_heads", 8),
-        "n_kv_heads": cfg.get("n_kv_heads", 2),
-        "d_ff": cfg.get("d_ff", 1536),
-        "max_len": cfg.get("max_len", 1028),
-        "rope_base": cfg.get("rope_base", 10000),
-        "dropout": 0.0,
-        "use_engram": cfg.get("use_engram", True),
-        "engram_every": cfg.get("engram_every", 2),
-        "engram_bucket_count": cfg.get("engram_bucket_count", 4096),
-        "engram_dim": cfg.get("engram_dim", 96),
-        "engram_order": cfg.get("engram_order", 3),
-        "pad_id": tokenizer.pad_token_id,
-    }
-    print(f"Model config: {resolved}")
-    model = CosmosT2_LLM(**resolved)
-    state = ckpt.get("model_state", ckpt)
-    state = {k.replace("module.", "", 1): v for k, v in state.items()}
-    missing, unexpected = model.load_state_dict(state, strict=False)
-    if missing:
-        print(f"Missing keys: {missing}")
-    if unexpected:
-        print(f"Unexpected keys: {unexpected}")
-    model = model.to(DEVICE).to(DTYPE).eval()
-    return model
-
-print(f"Device: {DEVICE} | dtype: {DTYPE}")
-tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token
-print(f"Tokenizer: {TOKENIZER_NAME}")
-model = load_checkpoint(tokenizer)
-n_params = sum(p.numel() for p in model.parameters())
-print(f"Loaded {n_params / 1e6:.2f}M parameters")
-EOS_ID = tokenizer.eos_token_id
-
-def build_prompt(history, user_msg, system_msg):
-    messages = []
-    if system_msg and system_msg.strip():
-        messages.append({"role": "system", "content": system_msg.strip()})
-    for item in history:
-        if isinstance(item, dict) and "role" in item and "content" in item:
-            messages.append({"role": item["role"], "content": item["content"]})
-    messages.append({"role": "user", "content": user_msg})
-    return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-
-def chat_fn(message, history, system_msg, temperature, top_k, ctx_size, max_new_tokens):
-    if not message or not message.strip():
-        yield ""
-        return
-    prompt = build_prompt(history, message, system_msg)
-    input_ids = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).input_ids.to(DEVICE)
-    ctx_size = int(min(max(int(ctx_size), 8), MAX_CTX_HARD))
-    if input_ids.shape[1] > ctx_size - 16:
-        input_ids = input_ids[:, -(ctx_size - 16):]
-    max_new_tokens = int(min(max(int(max_new_tokens), 1), MAX_NEW_HARD))
-    temperature = float(max(min(temperature, 2.0), 0.0))
-    top_k = int(max(top_k, 1))
-    cur_ids = input_ids
-    generated = []
-    partial_text = ""
-    yield partial_text
-    for _ in range(max_new_tokens):
-        nxt = model.generate_step(cur_ids, temperature=temperature, top_k=top_k, max_ctx=ctx_size)
-        cur_ids = torch.cat([cur_ids, nxt], dim=1)
-        token_id = int(nxt.item())
-        generated.append(token_id)
-        if token_id == EOS_ID:
-            break
-        new_text = tokenizer.decode(generated, skip_special_tokens=False)
-        if new_text != partial_text:
-            partial_text = new_text
-            yield partial_text
-    final_text = tokenizer.decode(generated, skip_special_tokens=False)
-    if final_text != partial_text:
-        yield final_text
-
-CUSTOM_CSS = """
-html, body { overflow-x: hidden; max-width: 100vw; }
-.gradio-container {
-    max-width: 1100px !important;
-    width: 100% !important;
-    margin: auto;
-    padding: 12px !important;
-    box-sizing: border-box;
-}
-.gradio-container *, .gradio-container *::before, .gradio-container *::after {
-    box-sizing: border-box;
-}
-.gradio-container .form, .gradio-container .block, .gradio-container .gap {
-    min-width: 0 !important;
-    max-width: 100% !important;
-}
-#header-card {
-    background: linear-gradient(135deg, #0d1117 0%, #161b22 100%);
-    border: 1px solid #30363d;
-    border-radius: 14px;
-    padding: 22px;
-    margin-bottom: 14px;
-    text-align: center;
-}
-#header-card h2 { color: #58a6ff; margin: 6px 0; font-weight: 600; }
-#header-card p { color: #8b949e; margin: 4px 0; font-size: 0.92em; }
-.badge {
-    display: inline-block;
-    background: #21262d;
-    color: #c9d1d9;
-    padding: 3px 10px;
-    border-radius: 999px;
-    font-size: 0.78em;
-    margin: 2px 4px;
-    border: 1px solid #30363d;
-}
-.warn {
-    background: #2a1f0a;
-    border: 1px solid #6b4d11;
-    color: #f0c674;
-    padding: 10px 14px;
-    border-radius: 10px;
-    font-size: 0.88em;
-    margin-top: 8px;
-    text-align: left;
-}
-.message-wrap, .message, .bubble, .gradio-container pre, .gradio-container code {
-    max-width: 100% !important;
-    overflow-wrap: break-word !important;
-    word-break: break-word !important;
-}
-.gradio-container pre {
-    white-space: pre-wrap !important;
-    overflow-x: auto;
-}
-footer { visibility: hidden; }
-@media (max-width: 640px) {
-    .gradio-container { padding: 6px !important; }
-    #header-card { padding: 14px 12px; border-radius: 10px; }
-    #header-card h2 { font-size: 1.05em; }
-    #header-card p { font-size: 0.85em; }
-    .badge { font-size: 0.7em; padding: 2px 8px; margin: 2px 2px; }
-    .warn { font-size: 0.8em; padding: 8px 10px; }
-    .gradio-container .chatbot, .gradio-container [data-testid="chatbot"] {
-        min-height: 380px !important;
-        max-height: 60vh !important;
-    }
-}
-"""
-
-HEADER_HTML = f"""
-<div id="header-card">
-  <h2>{MODEL_NAME}</h2>
-  <p>
-    <span class="badge">{n_params / 1e6:.2f}M params</span>
-    <span class="badge">{N_LAYERS} layers</span>
-    <span class="badge">RoPE + RMSNorm + SwiGLU + GQA</span>
-    <span class="badge">Engram {"on" if True else "off"}</span>
-  </p>
-  <p>
-    Trained from scratch on <a href="https://huggingface.co/datasets/$DATASET_NAME" target="_blank">$DATASET_NAME</a>
-    · Model repo: <a href="https://huggingface.co/wop/Cosmos-T2-80M-Test" target="_blank">wop/Cosmos-T2-80M-Test</a>
-  </p>
-  <div class="warn">
-    This is a research/demo model. It can overfit quickly and should not be treated as a factual assistant.
-  </div>
-</div>
-"""
-
-EXAMPLES = [
-    ["What is 12 * 7?", DEFAULT_SYSTEM_PROMPT, 0.1, 1, MAX_CTX_HARD, 128],
-    ["Write a haiku about debugging at 3 AM.", DEFAULT_SYSTEM_PROMPT, 0.1, 1, MAX_CTX_HARD, 128],
-    ["Why does ice float?", DEFAULT_SYSTEM_PROMPT, 0.1, 1, MAX_CTX_HARD, 128],
-]
-
-system_msg_box = gr.Textbox(label="System prompt", value=DEFAULT_SYSTEM_PROMPT, lines=3)
-temperature_slider = gr.Slider(minimum=0.0, maximum=2.0, value=0.1, step=0.05, label="Temperature")
-top_k_slider = gr.Slider(minimum=1, maximum=200, value=1, step=1, label="Top-K")
-ctx_size_slider = gr.Slider(minimum=64, maximum=MAX_CTX_HARD, value=MAX_CTX_HARD, step=64, label=f"Context window (max {MAX_CTX_HARD})")
-max_new_slider = gr.Slider(minimum=16, maximum=MAX_NEW_HARD, value=128, step=16, label=f"Max new tokens (max {MAX_NEW_HARD})")
-
-with gr.Blocks(
-    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="slate", neutral_hue="slate"),
-    css=CUSTOM_CSS,
-    title=f"{MODEL_NAME} Demo",
-) as demo:
-    gr.HTML(HEADER_HTML)
-    gr.ChatInterface(
-        fn=chat_fn,
-        type="messages",
-        additional_inputs=[system_msg_box, temperature_slider, top_k_slider, ctx_size_slider, max_new_slider],
-        additional_inputs_accordion=gr.Accordion(label="Generation Parameters", open=True),
-        examples=EXAMPLES,
-        cache_examples=False,
-        chatbot=gr.Chatbot(
-            type="messages",
-            height=520,
-            show_copy_button=True,
-            render_markdown=True,
-            sanitize_html=False,
-        ),
-        textbox=gr.Textbox(placeholder="Ask Cosmos-T2 anything…", autofocus=True, lines=1, max_lines=6),
-        title=None,
-        description=None,
-    )
-    gr.Markdown("**Tip:** keep temperature low for the most stable outputs. The model is research-grade and intentionally small.")
-
-if __name__ == "__main__":
-    demo.queue(max_size=16, default_concurrency_limit=1).launch(server_name="0.0.0.0", server_port=7860)