Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import os | |
| os.environ["KERAS_BACKEND"] = "jax" | |
| import numpy as np | |
| import jax | |
| import keras | |
| import gradio as gr | |
| import time | |
| from pathlib import Path | |
| from veylon_model import create_llm | |
| from tokenizer import TokenizerWrapper | |
| from config import ( | |
| CONTEXT, | |
| vocab_size, | |
| D_MODEL, | |
| numberoflayers, | |
| numberofheads, | |
| d_Latent, | |
| ffn_mult, | |
| num_kv_heads, | |
| swa_window, | |
| ) | |
| # ============================================================ | |
| # Initialize (runs once) | |
| # ============================================================ | |
| keras.mixed_precision.set_global_policy("mixed_bfloat16") | |
| print(f"Backend: {keras.backend.backend()}") | |
| print(f"JAX devices: {jax.devices()}") | |
| # Load tokenizer | |
| tokenizer = TokenizerWrapper("tokenizer.model") | |
| assert tokenizer.vocab_size == vocab_size, ( | |
| f"Tokenizer vocab ({tokenizer.vocab_size}) != config vocab ({vocab_size})" | |
| ) | |
| print(f"✓ Tokenizer loaded: {tokenizer.vocab_size} vocab") | |
| # Build model | |
| print("Building model...") | |
| model = create_llm( | |
| vocab_size=vocab_size, | |
| d_model=D_MODEL, | |
| n_layers=numberoflayers, | |
| n_heads=numberofheads, | |
| d_latent=d_Latent, | |
| ffn_mult=ffn_mult, | |
| max_seq_len=CONTEXT, | |
| use_moe=False, | |
| num_kv_heads=num_kv_heads, | |
| swa_window=swa_window, | |
| ) | |
| # Warmup | |
| dummy = np.zeros((1, CONTEXT), dtype=np.int32) | |
| _ = model(dummy, training=False) | |
| print("✓ Model built successfully") | |
| # Load weights | |
| WEIGHTS_PATH = "veylon_final.weights.h5" | |
| if Path(WEIGHTS_PATH).exists(): | |
| print(f"Loading weights from: {WEIGHTS_PATH}") | |
| model.load_weights(WEIGHTS_PATH) | |
| print("✓ Weights loaded successfully") | |
| else: | |
| print(f"WARNING: {WEIGHTS_PATH} not found. Using untrained model.") | |
| print(f"✓ Model params: {model.count_params():,}\n") | |
| # ============================================================ | |
| # Sampling | |
| # ============================================================ | |
| def sample_from_logits( | |
| logits_row: np.ndarray, | |
| temperature: float = 0.7, | |
| min_p: float = 0.05, | |
| ) -> int: | |
| """ | |
| Min-P sampling (replaces top-k). Top-k/top-p keep a FIXED-size or | |
| fixed-cumulative-mass candidate pool open regardless of how confident | |
| the model actually is -- on a small/uncertain model, that means when | |
| it doesn't know what comes next, top-k still hands it 50 candidates | |
| to choose from, many of them garbage. Min-P instead sets a threshold | |
| RELATIVE to the top token's probability: confident predictions | |
| collapse the pool to 1-2 tokens, uncertain predictions keep it wide. | |
| Expects and returns a 1D array (single vocab-length row), not (1, vocab). | |
| """ | |
| logits_row = np.array(logits_row, dtype=np.float32, copy=True) | |
| if temperature > 0: | |
| logits_row = logits_row / float(max(temperature, 1e-8)) | |
| row = logits_row - np.max(logits_row) # numerical stability | |
| probs = np.exp(row) | |
| probs = probs / probs.sum() | |
| if min_p > 0: | |
| threshold = min_p * probs.max() | |
| mask = probs >= threshold | |
| if not mask.any(): # degenerate guard -- never zero out everything | |
| mask[np.argmax(probs)] = True | |
| probs = np.where(mask, probs, 0.0) | |
| probs = probs / probs.sum() | |
| return int(np.random.choice(len(probs), p=probs)) | |
| def apply_repetition_controls( | |
| logits_row: np.ndarray, | |
| generated_ids: list, | |
| no_repeat_ngram_size: int = 3, | |
| repetition_penalty: float = 1.2, | |
| ) -> np.ndarray: | |
| """ | |
| Applied BEFORE sample_from_logits, on the model's own newly-generated | |
| tokens only (never the user's prompt -- penalizing someone for the | |
| model reusing words from THEIR prompt would be wrong, this is about | |
| stopping the model's own loops). | |
| 1. repetition_penalty (soft): divides a positive logit for an | |
| already-seen token, MULTIPLIES a negative one -- always pushes the | |
| token down, regardless of logit sign (a naive "divide by penalty" | |
| for a negative logit would actually push it UP, backwards). | |
| 2. no_repeat_ngram_size (hard): if continuing with a candidate token | |
| would recreate an n-gram that's already appeared in this | |
| generation, that candidate is set to -inf. This is what actually | |
| kills infinite loops -- it's a hard constraint, not a nudge. | |
| """ | |
| logits_row = logits_row.copy() | |
| if repetition_penalty and repetition_penalty != 1.0 and generated_ids: | |
| for tok_id in set(generated_ids): | |
| if logits_row[tok_id] > 0: | |
| logits_row[tok_id] /= repetition_penalty | |
| else: | |
| logits_row[tok_id] *= repetition_penalty | |
| n = no_repeat_ngram_size | |
| if n and n > 0 and len(generated_ids) >= n - 1: | |
| prefix = tuple(generated_ids[-(n - 1):]) if n > 1 else tuple() | |
| banned = set() | |
| for i in range(len(generated_ids) - n + 1): | |
| if tuple(generated_ids[i:i + n - 1]) == prefix: | |
| banned.add(generated_ids[i + n - 1]) | |
| for tok_id in banned: | |
| logits_row[tok_id] = -np.inf | |
| return logits_row | |
| # Few-shot "anchoring": Veylon is a CONTINUATION-only base model, never | |
| # instruction-tuned -- it has no concept of "Input:"/"Output:" turns, so | |
| # priming it with instruction-style examples (as generic small-LM advice | |
| # often suggests) would push it OUT of its training distribution, not | |
| # into it. What anchoring means correctly here is different: prepend a | |
| # short, clean, well-formed piece of narrative prose so the model's | |
| # recent-attention window is dominated by coherent English *before* | |
| # generation starts, nudging it away from degenerating into word salad. | |
| # This consumes some of the context budget, so it's short and toggleable. | |
| ANCHOR_TEXT = ( | |
| "The old lighthouse stood at the edge of the cliff, its light " | |
| "sweeping slowly across the dark water. Every night, the keeper " | |
| "climbed the narrow stairs to check the lamp, and every night the " | |
| "sea answered with the steady sound of waves against the rocks.\n\n" | |
| ) | |
| # ============================================================ | |
| # Streaming generation with live telemetry | |
| # ============================================================ | |
| def respond( | |
| message: str, | |
| history: list, | |
| max_new_tokens: int = 96, | |
| temperature: float = 0.7, | |
| min_p: float = 0.05, | |
| no_repeat_ngram_size: int = 3, | |
| repetition_penalty: float = 1.2, | |
| use_anchor: bool = True, | |
| ): | |
| """ | |
| gr.ChatInterface-compatible streaming generator. Yields incrementally | |
| as tokens are produced, and appends a real (measured, not decorative) | |
| telemetry line -- token count / time-to-first-token / tokens-per- | |
| second -- under the reply, matching the "signature element" of this | |
| UI: honest, live numbers about what the model is actually doing. | |
| IMPORTANT, honest by design: `history` is accepted (gr.ChatInterface | |
| requires the signature) and IS shown to the person as an ongoing | |
| conversation, but is deliberately NOT concatenated into the prompt | |
| fed to the model. Arya is a continuation-only base model, never | |
| instruction-tuned on multi-turn User/Assistant-style dialogue at | |
| scale -- feeding it a growing instruction-formatted transcript would | |
| push it further out of its training distribution with every turn, | |
| not closer to coherence. Each message is generated as its own fresh | |
| continuation (with the anchor primer). This is stated plainly in the | |
| UI caption rather than left as a surprise. | |
| """ | |
| try: | |
| full_prompt = (ANCHOR_TEXT + message) if use_anchor else message | |
| tokens = tokenizer.encode(full_prompt, add_bos=True, add_eos=False) | |
| if len(tokens) == 0: | |
| tokens = [tokenizer.bos_id if hasattr(tokenizer, "bos_id") else 1] | |
| tokens = tokens[-CONTEXT:] | |
| generated_ids: list = [] | |
| t_start = time.perf_counter() | |
| ttft = None | |
| prompt_ids = np.array([tokens], dtype=np.int32) | |
| logits, cache_k, cache_v = model.generate_step( | |
| prompt_ids, cache_k=None, cache_v=None, cache_pos=0, | |
| ) | |
| logits_row = np.array(logits[0, -1, :], dtype=np.float32, copy=True) | |
| logits_row = apply_repetition_controls( | |
| logits_row, generated_ids, no_repeat_ngram_size, repetition_penalty | |
| ) | |
| next_token = sample_from_logits(logits_row, temperature=temperature, min_p=min_p) | |
| tokens.append(next_token) | |
| generated_ids.append(next_token) | |
| ttft = time.perf_counter() - t_start | |
| partial_text = tokenizer.decode(generated_ids) | |
| yield _with_telemetry(partial_text, len(generated_ids), ttft, t_start) | |
| if next_token != tokenizer.eos_id and len(tokens) < CONTEXT: | |
| cache_pos = len(prompt_ids[0]) | |
| for _ in range(max_new_tokens - 1): | |
| next_input = np.array([[next_token]], dtype=np.int32) | |
| logits, cache_k, cache_v = model.generate_step( | |
| next_input, cache_k=cache_k, cache_v=cache_v, cache_pos=cache_pos, | |
| ) | |
| cache_pos += 1 | |
| logits_row = np.array(logits[0, -1, :], dtype=np.float32, copy=True) | |
| logits_row = apply_repetition_controls( | |
| logits_row, generated_ids, no_repeat_ngram_size, repetition_penalty | |
| ) | |
| next_token = sample_from_logits(logits_row, temperature=temperature, min_p=min_p) | |
| tokens.append(next_token) | |
| generated_ids.append(next_token) | |
| if next_token == tokenizer.eos_id: | |
| break | |
| partial_text = tokenizer.decode(generated_ids) | |
| yield _with_telemetry(partial_text, len(generated_ids), ttft, t_start) | |
| if len(tokens) >= CONTEXT: | |
| break | |
| except Exception as e: | |
| yield f"Error: {str(e)}" | |
| def _with_telemetry(text: str, n_tok: int, ttft: float, t_start: float) -> str: | |
| elapsed = max(time.perf_counter() - t_start, 1e-6) | |
| tok_per_s = n_tok / elapsed | |
| telemetry = ( | |
| f"\n\n<sub style='font-family:monospace;color:var(--body-text-color-subdued)'>" | |
| f"{n_tok} tok · TTFT {ttft * 1000:.0f} ms · {tok_per_s:.1f} tok/s</sub>" | |
| ) | |
| return text + telemetry | |
| # ============================================================ | |
| # Gradio UI — Arya | |
| # ============================================================ | |
| # NOTE on "gr.server()": there is no such function in Gradio. gr.Blocks | |
| # and gr.ChatInterface are ALREADY built on FastAPI/Uvicorn internally -- | |
| # demo.launch() below is what actually starts that server. Same | |
| # clarification as the earlier "use FastAPI for speed" question: the web | |
| # framework was never the bottleneck for a custom JAX model like this. | |
| ARYA_CSS = """ | |
| :root { | |
| --arya-bg: #100C1A; | |
| --arya-bg-2: #17111F; | |
| --arya-panel: #1D1730; | |
| --arya-panel-2: #271F42; | |
| --arya-accent: #E8A33D; | |
| --arya-accent-soft: #F0BE73; | |
| --arya-accent-2: #5FBFA8; | |
| --arya-indigo: #5B67D9; | |
| --arya-text: #F3EEE3; | |
| --arya-text-dim: #948BAA; | |
| --arya-hairline: #2A2440; | |
| /* Retheme Gradio's OWN documented CSS variables (verified against the | |
| installed theme's actual variable names) instead of guessing at | |
| hashed/scoped internal component classes -- this cascades correctly | |
| into sliders, checkboxes, accordions, buttons, etc. automatically, | |
| and won't silently stop working on a Gradio version bump the way | |
| hand-guessed internal class names would. */ | |
| --body-background-fill: var(--arya-bg); | |
| --background-fill-primary: var(--arya-bg); | |
| --background-fill-secondary: var(--arya-panel); | |
| --block-background-fill: var(--arya-panel); | |
| --block-border-color: var(--arya-hairline); | |
| --block-title-text-color: var(--arya-text-dim); | |
| --block-label-text-color: var(--arya-text-dim); | |
| --body-text-color: var(--arya-text); | |
| --body-text-color-subdued: var(--arya-text-dim); | |
| --border-color-primary: var(--arya-hairline); | |
| --border-color-accent: var(--arya-accent); | |
| --color-accent: var(--arya-accent); | |
| --color-accent-soft: color-mix(in srgb, var(--arya-accent) 20%, transparent); | |
| --input-background-fill: var(--arya-panel-2); | |
| --input-border-color: var(--arya-hairline); | |
| --input-border-color-focus: var(--arya-accent); | |
| --input-placeholder-color: var(--arya-text-dim); | |
| --slider-color: var(--arya-accent); | |
| --checkbox-background-color-selected: var(--arya-accent); | |
| --checkbox-border-color-selected: var(--arya-accent); | |
| --checkbox-label-background-fill-selected: color-mix(in srgb, var(--arya-accent) 18%, var(--arya-panel)); | |
| --button-primary-background-fill: linear-gradient(135deg, var(--arya-accent), var(--arya-accent-soft)); | |
| --button-primary-background-fill-hover: var(--arya-accent-soft); | |
| --button-primary-text-color: #201306; | |
| --button-secondary-background-fill: var(--arya-panel-2); | |
| --button-secondary-border-color: var(--arya-hairline); | |
| --button-secondary-background-fill-hover: var(--arya-panel); | |
| --link-text-color: var(--arya-accent-2); | |
| } | |
| * { scrollbar-width: thin; scrollbar-color: var(--arya-panel-2) transparent; } | |
| ::-webkit-scrollbar { width: 8px; height: 8px; } | |
| ::-webkit-scrollbar-thumb { background: var(--arya-panel-2); border-radius: 8px; } | |
| ::-webkit-scrollbar-thumb:hover { background: var(--arya-hairline); } | |
| .gradio-container { | |
| background: | |
| radial-gradient(ellipse 900px 400px at 15% -10%, color-mix(in srgb, var(--arya-accent) 10%, transparent), transparent 60%), | |
| radial-gradient(ellipse 700px 400px at 100% 0%, color-mix(in srgb, var(--arya-indigo) 12%, transparent), transparent 55%), | |
| var(--arya-bg) !important; | |
| font-family: 'Inter', system-ui, sans-serif !important; | |
| max-width: 880px !important; | |
| } | |
| /* ── Topbar ── */ | |
| #arya-topbar { | |
| display: flex; | |
| align-items: center; | |
| gap: 12px; | |
| padding: 18px 4px 16px 4px; | |
| border-bottom: 1px solid var(--arya-hairline); | |
| margin-bottom: 10px; | |
| } | |
| #arya-mark { | |
| width: 34px; height: 34px; flex-shrink: 0; | |
| display: flex; align-items: center; justify-content: center; | |
| border-radius: 10px; | |
| background: linear-gradient(135deg, var(--arya-accent), var(--arya-indigo)); | |
| box-shadow: 0 2px 14px color-mix(in srgb, var(--arya-accent) 35%, transparent); | |
| } | |
| #arya-mark svg { width: 19px; height: 19px; } | |
| #arya-title-block { display: flex; flex-direction: column; line-height: 1.15; } | |
| #arya-title { | |
| font-family: 'IBM Plex Mono', monospace; | |
| font-weight: 700; | |
| font-size: 1.18rem; | |
| letter-spacing: 0.14em; | |
| color: var(--arya-text); | |
| } | |
| #arya-subtitle { | |
| font-family: 'IBM Plex Mono', monospace; | |
| font-size: 0.66rem; | |
| letter-spacing: 0.06em; | |
| color: var(--arya-text-dim); | |
| } | |
| #arya-status-chip { | |
| font-family: 'IBM Plex Mono', monospace; | |
| font-size: 0.72rem; | |
| font-weight: 500; | |
| color: var(--arya-accent-2); | |
| display: flex; align-items: center; gap: 6px; | |
| padding: 5px 10px; | |
| border-radius: 999px; | |
| background: color-mix(in srgb, var(--arya-accent-2) 12%, transparent); | |
| border: 1px solid color-mix(in srgb, var(--arya-accent-2) 30%, transparent); | |
| margin-left: 14px; | |
| } | |
| #arya-status-chip::before { | |
| content: ""; | |
| width: 6px; height: 6px; border-radius: 50%; | |
| background: var(--arya-accent-2); | |
| box-shadow: 0 0 8px var(--arya-accent-2); | |
| animation: arya-pulse 2s ease-in-out infinite; | |
| } | |
| @keyframes arya-pulse { | |
| 0%, 100% { opacity: 1; transform: scale(1); } | |
| 50% { opacity: 0.5; transform: scale(0.8); } | |
| } | |
| #arya-specs { | |
| margin-left: auto; | |
| display: flex; gap: 6px; flex-wrap: wrap; justify-content: flex-end; | |
| } | |
| .arya-chip { | |
| font-family: 'IBM Plex Mono', monospace; | |
| font-size: 0.66rem; | |
| letter-spacing: 0.03em; | |
| color: var(--arya-text-dim); | |
| padding: 4px 9px; | |
| border-radius: 6px; | |
| background: var(--arya-panel); | |
| border: 1px solid var(--arya-hairline); | |
| white-space: nowrap; | |
| } | |
| /* ── Chat transcript (verified real Gradio 6 Chatbot classes, not guessed: | |
| .user-row / .bot-row / .message-bubble-border confirmed present in the | |
| installed package's compiled component CSS) ── */ | |
| .user-row { | |
| background: linear-gradient(135deg, var(--arya-accent), var(--arya-accent-soft)) !important; | |
| color: #201306 !important; | |
| } | |
| .user-row * { color: #201306 !important; } | |
| .bot-row { | |
| background: var(--arya-panel) !important; | |
| border: 1px solid var(--arya-hairline) !important; | |
| } | |
| .message-bubble-border { | |
| border-color: var(--arya-hairline) !important; | |
| box-shadow: 0 4px 18px rgba(0,0,0,0.25); | |
| animation: arya-rise 0.25s ease-out; | |
| } | |
| @keyframes arya-rise { | |
| from { opacity: 0; transform: translateY(6px); } | |
| to { opacity: 1; transform: translateY(0); } | |
| } | |
| .bot-row code, .bot-row pre { | |
| font-family: 'IBM Plex Mono', monospace !important; | |
| background: var(--arya-bg-2) !important; | |
| border: 1px solid var(--arya-hairline) !important; | |
| border-radius: 6px !important; | |
| } | |
| /* ── Input row ── */ | |
| textarea, input[type="text"] { | |
| transition: border-color 0.15s ease, box-shadow 0.15s ease !important; | |
| } | |
| textarea:focus, input[type="text"]:focus { | |
| box-shadow: 0 0 0 3px color-mix(in srgb, var(--arya-accent) 22%, transparent) !important; | |
| } | |
| button.primary { | |
| transition: transform 0.12s ease, box-shadow 0.12s ease !important; | |
| box-shadow: 0 2px 12px color-mix(in srgb, var(--arya-accent) 30%, transparent) !important; | |
| } | |
| button.primary:hover { transform: translateY(-1px) scale(1.03); } | |
| /* ── Settings accordion & example chips: themed via the CSS variables | |
| above, no internal class names needed ── */ | |
| .example { | |
| transition: transform 0.12s ease, border-color 0.12s ease !important; | |
| } | |
| .example:hover { | |
| transform: translateY(-2px); | |
| border-color: var(--arya-accent) !important; | |
| } | |
| #arya-caption { | |
| font-family: 'IBM Plex Mono', monospace; | |
| font-size: 0.68rem; | |
| color: var(--arya-text-dim); | |
| text-align: center; | |
| padding: 16px 20px 6px 20px; | |
| letter-spacing: 0.03em; | |
| border-top: 1px solid var(--arya-hairline); | |
| margin-top: 12px; | |
| } | |
| """ | |
| ARYA_HEAD = """ | |
| <link rel="preconnect" href="https://fonts.googleapis.com"> | |
| <link href="https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&family=IBM+Plex+Mono:wght@400;500;700&display=swap" rel="stylesheet"> | |
| """ | |
| # Simple geometric monogram, not a decorative flourish -- the gradient fill | |
| # (accent -> indigo) matches #arya-mark's CSS background, single mark | |
| # rather than a busier logo, keeps it legible at 19px. | |
| ARYA_MARK_SVG = """ | |
| <svg viewBox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg"> | |
| <path d="M12 3L21 20H16.5L12 11L7.5 20H3L12 3Z" fill="#100C1A"/> | |
| </svg> | |
| """ | |
| def main(): | |
| spec_chips = [ | |
| f"{model.count_params() / 1e6:.1f}M PARAMS", | |
| f"{numberoflayers}L", | |
| f"CTX {CONTEXT}", | |
| f"GQA {numberofheads}/{num_kv_heads}", | |
| ] | |
| spec_html = "".join(f'<span class="arya-chip">{c}</span>' for c in spec_chips) | |
| with gr.Blocks(title="Arya") as demo: | |
| gr.HTML(f""" | |
| <div id="arya-topbar"> | |
| <div id="arya-mark">{ARYA_MARK_SVG}</div> | |
| <div id="arya-title-block"> | |
| <span id="arya-title">ARYA</span> | |
| <span id="arya-subtitle">FROM-SCRATCH SMALL LM</span> | |
| </div> | |
| <span id="arya-status-chip">READY</span> | |
| <div id="arya-specs">{spec_html}</div> | |
| </div> | |
| """) | |
| gr.ChatInterface( | |
| fn=respond, | |
| additional_inputs=[ | |
| gr.Slider(10, 256, value=96, step=10, label="Max tokens"), | |
| gr.Slider(0.1, 2.0, value=0.7, step=0.1, label="Temperature"), | |
| gr.Slider(0.0, 0.3, value=0.05, step=0.01, label="Min-P", | |
| info="Replaces top-k. Higher = stricter."), | |
| gr.Slider(0, 6, value=3, step=1, label="No-repeat n-gram size", | |
| info="Hard-blocks repeating any phrase this long. 0 disables."), | |
| gr.Slider(1.0, 1.5, value=1.2, step=0.05, label="Repetition penalty"), | |
| gr.Checkbox(value=True, label="Narrative anchor", | |
| info="Primes the model with clean prose before your prompt."), | |
| ], | |
| additional_inputs_accordion=gr.Accordion("Settings", open=False), | |
| examples=[ | |
| ["Once upon a time, in a village by the mountains,"], | |
| ["The scientist looked at the data and realized"], | |
| ["Explain photosynthesis simply."], | |
| ], | |
| chatbot=gr.Chatbot(height=460, show_label=False), | |
| textbox=gr.Textbox(placeholder="Ask Arya anything...", show_label=False), | |
| ) | |
| gr.HTML(""" | |
| <div id="arya-caption"> | |
| CONTINUATION MODEL, NOT INSTRUCTION-TUNED — EACH REPLY IS A FRESH | |
| GENERATION, NOT A REMEMBERED CONVERSATION · RUNS FULLY ON YOUR HARDWARE | |
| </div> | |
| """) | |
| # css/head/theme live on launch(), not Blocks(), as of Gradio 6 -- | |
| # verified against the actually-installed version rather than assumed. | |
| demo.launch( | |
| share=True, server_name="0.0.0.0", server_port=7860, | |
| css=ARYA_CSS, head=ARYA_HEAD, | |
| theme=gr.themes.Base(primary_hue=gr.themes.colors.orange, | |
| neutral_hue=gr.themes.colors.slate), | |
| ) | |
| if __name__ == "__main__": | |
| main() |