Delete inference.py

inference.py

@@ -1,193 +0,0 @@
-from __future__ import annotations
-
-import os
-os.environ["KERAS_BACKEND"] = "jax"
-
-import numpy as np
-import jax
-import keras
-
-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,
-)
-
-# ============================================================
-# Runtime info
-# ============================================================
-
-print(f"Backend: {keras.backend.backend()}")
-print(f"JAX devices: {jax.devices()}")
-
-keras.mixed_precision.set_global_policy("mixed_bfloat16")
-
-# ============================================================
-# Load tokenizer
-# ============================================================
-
-tokenizer = TokenizerWrapper("tokenizer.model")
-
-assert tokenizer.vocab_size == vocab_size, (
-    f"Tokenizer vocab ({tokenizer.vocab_size}) "
-    f"!= config vocab ({vocab_size})"
-)
-
-print(f"Tokenizer vocab size: {tokenizer.vocab_size}")
-
-# ============================================================
-# Build model (must exactly match training)
-# ============================================================
-
-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 with EXACT training/inference shape
-dummy = np.zeros((1, CONTEXT), dtype=np.int32)
-_ = model(dummy, training=False)
-
-print("✓ Model built successfully")
-
-# ============================================================
-# Load weights
-# ============================================================
-
-WEIGHTS_PATH = "veylon_final.weights.h5"
-print(f"Loading weights from: {WEIGHTS_PATH}")
-model.load_weights(WEIGHTS_PATH)
-print("✓ Weights loaded successfully")
-
-# ============================================================
-# Sampling settings
-# ============================================================
-
-MAX_NEW_TOKENS = 64
-TEMPERATURE = 0.8
-TOP_K = 50
-
-def sample_from_logits(
-    logits: np.ndarray,
-    temperature: float = 0.8,
-    top_k: int = 50,
-) -> int:
-    """
-    NumPy-only sampling to avoid JAX/readonly array issues.
-    """
-    logits = np.array(logits, dtype=np.float32, copy=True)
-
-    if temperature > 0:
-        logits = logits / float(max(temperature, 1e-8))
-
-    if top_k > 0:
-        k = min(int(top_k), logits.shape[-1])
-        row = logits[0]
-        top_indices = np.argpartition(row, -k)[-k:]
-
-        filtered = np.full_like(row, -np.inf)
-        filtered[top_indices] = row[top_indices]
-        logits[0] = filtered
-
-    row = logits[0]
-    row = row - np.max(row)
-    probs = np.exp(row)
-    probs = probs / probs.sum()
-
-    return int(np.random.choice(len(probs), p=probs))
-
-# ============================================================
-# Generation loop
-# ============================================================
-
-while True:
-    prompt = input("\nEnter your prompt (or 'exit'): ").strip()
-
-    if prompt.lower() in {"exit", "quit"}:
-        break
-
-    tokens = tokenizer.encode(
-        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:]
-
-    print("\nGenerating...\n")
-
-    # Prompt prefill (one-time)
-    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,
-    )
-
-    next_token = sample_from_logits(
-        np.array(logits[:, -1, :], dtype=np.float32, copy=True),
-        temperature=TEMPERATURE,
-        top_k=TOP_K,
-    )
-    tokens.append(next_token)
-
-    if next_token != tokenizer.eos_id and len(tokens) < CONTEXT:
-        # After prefill, we are decoding token-by-token.
-        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
-
-            next_token = sample_from_logits(
-                np.array(logits[:, -1, :], dtype=np.float32, copy=True),
-                temperature=TEMPERATURE,
-                top_k=TOP_K,
-            )
-            tokens.append(next_token)
-
-            if next_token == tokenizer.eos_id:
-                break
-
-            if len(tokens) >= CONTEXT:
-                print("\n[Context limit reached]")
-                break
-
-    generated_text = tokenizer.decode(tokens)
-
-    print("\n" + "=" * 60)
-    print("Veylon Alpha")
-    print("=" * 60)
-    print(generated_text)
-    print("=" * 60)