Update app.py

app.py

@@ -1,5 +1,4 @@
-# app.py — CodVa-2 Demo with streaming
-
+# app.py — CodVa-2 Demo (PRETRAIN model, domain tokens only)
 import os
 import math
 import time
@@ -10,14 +9,17 @@ import gradio as gr
 from dataclasses import dataclass
 from typing import Tuple, Generator
 from tokenizers import Tokenizer
-from huggingface_hub import hf_hub_download, login, HfApi  # login not hf_login
+from huggingface_hub import hf_hub_download, login, HfApi
 
+# ─────────────────────────────────────────────────────────────────────────────
+#  AUTH
+# ─────────────────────────────────────────────────────────────────────────────
 HF_TOKEN        = os.environ.get("HF_TOKEN", "")
 HF_DATASET_REPO = "Bc-AI/nova1_data"
 HF_MODEL_REPO   = os.environ.get("MODEL_REPO", "hugging-science/CodVa-2-session-002")
 
 if HF_TOKEN:
-    login(token=HF_TOKEN)  # login not hf_login
+    login(token=HF_TOKEN)
 
 # ─────────────────────────────────────────────────────────────────────────────
 #  CONFIG
@@ -241,11 +243,16 @@ tok_path  = hf_hub_download(
     token=HF_TOKEN or None,
 )
 tokenizer = Tokenizer.from_file(tok_path)
-EOS_ID    = (
-    tokenizer.token_to_id("<|endoftext|>")
-    or tokenizer.token_to_id("</s>")
-    or -1
-)
+
+# Domain tokens
+DOMAIN_TOKENS = {
+    "Code":      "<|domain_code|>",
+    "Math":      "<|domain_math|>",
+    "General":   "<|domain_general|>",
+    "Reasoning": "<|domain_reasoning|>",
+}
+
+EOS_ID = tokenizer.token_to_id("<|endoftext|>") or tokenizer.token_to_id("</s>") or -1
 print(f"[init] tokenizer | vocab={tokenizer.get_vocab_size()} | eos={EOS_ID}")
 
 print("[init] loading model...")
@@ -281,11 +288,8 @@ else:
     print("[init] WARNING: no weights found — random init")
 
 model.eval()
-# Compile for ~30% CPU speedup (optional, remove if it causes issues)
-print("[init] model ready (CPU inference, no compile)")
-
 n_params = sum(p.numel() for p in model.parameters())
-print(f"[init] ready | {n_params/1e6:.1f}M params")
+print(f"[init] ready | {n_params/1e6:.1f}M params | CPU inference")
 
 # ─────────────────────────────────────────────────────────────────────────────
 #  STREAMING GENERATION
@@ -293,22 +297,29 @@ print(f"[init] ready | {n_params/1e6:.1f}M params")
 @torch.no_grad()
 def generate_stream(
     prompt:      str,
+    domain:      str,
     max_new:     int   = 256,
     temperature: float = 0.8,
     top_p:       float = 0.95,
     top_k:       int   = 50,
 ) -> Generator[Tuple[str, str], None, None]:
     """
-    Yields (full_text_so_far, stats_string) after every token.
-    Gradio detects the generator and streams automatically.
+    Pretrain-style generation with domain token prepending.
+    Model sees: <|domain_X|>{prompt}
+    Continues from there.
     """
     if not prompt or not prompt.strip():
         yield "", "⚠️ Please enter a prompt"
         return
 
-    enc = tokenizer.encode(prompt)
+    # Prepend domain token (matches training data format)
+    domain_token = DOMAIN_TOKENS.get(domain, "<|domain_code|>")
+    full_prompt  = domain_token + prompt.strip()
+
+    enc = tokenizer.encode(full_prompt)
     ids = enc.ids
-    # Truncate prompt if too long
+    
+    # Truncate if too long
     if len(ids) > cfg.max_len - max_new:
         ids = ids[-(cfg.max_len - max_new):]
 
@@ -343,32 +354,40 @@ def generate_stream(
         generated.append(next_tok)
         x = torch.cat([x, torch.tensor([[next_tok]])], dim=1)
 
-        # Decode full output so far
-        # Full decode is safer than incremental for BPE tokenizers
-        out_text = tokenizer.decode(generated, skip_special_tokens=False)
+        # Decode what we have so far (strip domain token from display)
+        full_text = tokenizer.decode(ids + generated, skip_special_tokens=False)
+        
+        # Remove domain token from display
+        display_text = full_text
+        for tok in DOMAIN_TOKENS.values():
+            display_text = display_text.replace(tok, "")
 
-        # Build stats string
+        # Build stats
         elapsed = time.time() - t0
         tps     = len(generated) / max(elapsed, 1e-3)
         stats   = (
             f"⏱ {elapsed:.1f}s  |  "
             f"🔤 {len(generated)} / {max_new} tokens  |  "
             f"⚡ {tps:.1f} tok/s  |  "
+            f"🎯 {domain}  |  "
             f"🌡 {temperature}  top-p {top_p}  top-k {int(top_k)}"
         )
 
-        yield prompt + out_text, stats
+        yield display_text, stats
 
         # Stop on EOS
         if EOS_ID >= 0 and next_tok == EOS_ID:
             break
 
-    # Final yield with done marker
-    out_text = tokenizer.decode(generated, skip_special_tokens=False)
-    elapsed  = time.time() - t0
-    tps      = len(generated) / max(elapsed, 1e-3)
+    # Final yield
+    full_text = tokenizer.decode(ids + generated, skip_special_tokens=False)
+    for tok in DOMAIN_TOKENS.values():
+        full_text = full_text.replace(tok, "")
+    
+    elapsed = time.time() - t0
+    tps     = len(generated) / max(elapsed, 1e-3)
     yield (
-        prompt + out_text,
+        full_text,
         f"✅ Done  |  ⏱ {elapsed:.1f}s  |  "
         f"🔤 {len(generated)} tokens  |  ⚡ {tps:.1f} tok/s"
     )
@@ -377,46 +396,59 @@ def generate_stream(
 #  GRADIO UI
 # ─────────────────────────────────────────────────────────────────────────────
 EXAMPLES = [
-    ["def fibonacci(n):\n    ",                                                    128, 0.2, 0.95, 50],
-    ["class BinarySearchTree:\n    def __init__(self):\n        ",                 256, 0.4, 0.95, 50],
-    ["# Sort a list using quicksort\ndef quicksort(arr):\n    ",                   200, 0.3, 0.95, 50],
-    ["import torch\nimport torch.nn as nn\n\nclass Attention(nn.Module):\n    ",   300, 0.5, 0.95, 50],
-    ["SELECT users.name, orders.total FROM users JOIN orders ON ",                 128, 0.4, 0.90, 40],
-    ["def binary_search(arr, target):\n    left, right = 0, len(arr) - 1\n    ",  150, 0.2, 0.95, 50],
-    ["// React component for a todo list\nfunction TodoApp() {\n  ",               200, 0.6, 0.95, 50],
-    ["fn quicksort(arr: &mut Vec<i32>) {\n    ",                                   200, 0.3, 0.95, 50],
+    ["def fibonacci(n):\n    ", "Code", 128, 0.2, 0.95, 50],
+    ["class BinaryTree:\n    def __init__(self):\n        ", "Code", 256, 0.3, 0.95, 50],
+    ["import torch\nimport torch.nn as nn\n\n", "Code", 200, 0.4, 0.95, 50],
+    ["SELECT users.name, orders.total FROM ", "Code", 100, 0.3, 0.90, 40],
+    ["# Quicksort implementation\ndef quicksort(arr):\n    ", "Code", 200, 0.2, 0.95, 50],
+    ["Problem: Find the derivative of f(x) = x^3 + 2x^2 - 5x + 1\n\nSolution: ", "Math", 150, 0.4, 0.95, 50],
+    ["Theorem: The sum of angles in a triangle equals 180 degrees.\n\nProof: ", "Math", 200, 0.5, 0.95, 50],
+    ["Let $f(x) = \\int_0^x t^2 dt$. Then ", "Math", 128, 0.3, 0.95, 50],
+    ["The history of the Roman Empire began ", "General", 200, 0.7, 0.95, 50],
+    ["Photosynthesis is the process by which ", "General", 150, 0.5, 0.95, 50],
 ]
 
 CSS = """
-.container   { max-width: 960px; margin: auto; }
-.code-box    { font-family: 'JetBrains Mono', 'Fira Code', 'Courier New', monospace !important;
-               font-size: 13px !important; line-height: 1.5 !important; }
-.stats-text  { font-size: 11px !important; color: #888 !important; }
-footer       { display: none !important; }
+.container { max-width: 1100px; margin: auto; }
+.code-box { 
+    font-family: 'JetBrains Mono', 'Fira Code', 'Courier New', monospace !important;
+    font-size: 13px !important; 
+    line-height: 1.5 !important; 
+}
 """
 
-with gr.Blocks(title="CodVa-2 Demo") as demo:
+with gr.Blocks(title="CodVa-2 Pretrain Demo") as demo:
 
+    gr.HTML(f"<style>{CSS}</style>")
+    
     gr.Markdown("""
-    # 🧠 CodVa-2 — Code Language Model
-    **213M parameters** · Trained on 9.24B tokens of code · PPL ~7.9 on validation
-
-    Architecture highlights:
-    **Differential Attention** · **Local/Global alternating layers** · **Dual-path FFN** · **Token Importance Gate**
+    # 🧠 CodVa-2 — Pretrained Code LM
+    **213M parameters** · Differential Attention · Trained on code/math/general corpus
+    
+    This is a **pretrained** model (not instruction-tuned). It continues text in the style of its training domain.
+    Use the domain selector to control what kind of continuation you get.
     """)
 
     with gr.Row():
         # ── Left: inputs ──────────────────────────────────────────────────────
         with gr.Column(scale=1):
             prompt_box = gr.Textbox(
-                label="Prompt",
-                placeholder="Start your code here...",
+                label="Prompt (raw text, model will continue)",
+                placeholder="def fibonacci(n):\n    ",
                 lines=10,
                 elem_classes=["code-box"],
             )
+            
+            domain_dropdown = gr.Dropdown(
+                choices=list(DOMAIN_TOKENS.keys()),
+                value="Code",
+                label="Domain (prepends domain token)",
+                info="Code, Math, General, or Reasoning — tells the model what style to use"
+            )
+            
             with gr.Row():
                 max_new_slider = gr.Slider(16, 512, value=256, step=16,
-                                           label="Max tokens")
+                                           label="Max new tokens")
                 temp_slider    = gr.Slider(0.0, 2.0, value=0.8, step=0.05,
                                            label="Temperature")
             with gr.Row():
@@ -432,7 +464,7 @@ with gr.Blocks(title="CodVa-2 Demo") as demo:
         # ── Right: output ─────────────────────────────────────────────────────
         with gr.Column(scale=1):
             output_box = gr.Textbox(
-                label="Output (streaming)",
+                label="Generated continuation (streaming)",
                 lines=20,
                 interactive=False,
                 elem_classes=["code-box"],
@@ -441,36 +473,38 @@ with gr.Blocks(title="CodVa-2 Demo") as demo:
                 label="",
                 lines=1,
                 interactive=False,
-                elem_classes=["stats-text"],
             )
 
     gr.Examples(
         examples=EXAMPLES,
-        inputs=[prompt_box, max_new_slider, temp_slider, topp_slider, topk_slider],
-        label="📋 Examples — click to load",
-        examples_per_page=8,
+        inputs=[prompt_box, domain_dropdown, max_new_slider, temp_slider, topp_slider, topk_slider],
+        label="📋 Example prompts — click to load",
+        examples_per_page=10,
     )
 
     gr.Markdown("""
     ---
     💡 **Tips:**  
-    `temp 0.1–0.3` = precise completions · `temp 0.7–1.2` = creative exploration  
-    Lower `top-p` = more focused · Higher `top-k` = more diverse vocabulary
+    - **Domain matters:** Code domain → code syntax, Math → equations, General → prose  
+    - **Lower temp (0.1-0.3)** = deterministic, predictable (good for code)  
+    - **Higher temp (0.7-1.2)** = creative, varied (good for text)  
+    - This model has seen **~2B tokens** (20% trained). Expect coherent syntax but sometimes wrong logic.  
+    - By 10B tokens it should be much stronger.
     """)
 
     # ── Wire up events ────────────────────────────────────────────────────────
     gen_event = gen_btn.click(
         fn=generate_stream,
-        inputs=[prompt_box, max_new_slider, temp_slider, topp_slider, topk_slider],
+        inputs=[prompt_box, domain_dropdown, max_new_slider, temp_slider, topp_slider, topk_slider],
         outputs=[output_box, stats_box],
     )
-    # Also trigger on Enter in the prompt box
+    
     prompt_box.submit(
         fn=generate_stream,
-        inputs=[prompt_box, max_new_slider, temp_slider, topp_slider, topk_slider],
+        inputs=[prompt_box, domain_dropdown, max_new_slider, temp_slider, topp_slider, topk_slider],
         outputs=[output_box, stats_box],
     )
-    # Stop button cancels the running generator
+    
     stop_btn.click(fn=None, cancels=[gen_event])
 
     clear_btn.click(
@@ -483,5 +517,4 @@ if __name__ == "__main__":
         server_name="0.0.0.0",
         server_port=7860,
         show_error=True,
-        css=CSS,   # <-- moved here
     )