app.py

# app.py — CodVa-2 SFT CPU backend
# Single file, KV cache, SDPA, optional INT8 quantization.

import os

# Set before importing torch.
CPU_THREADS = int(
    os.environ.get(
        "CPU_THREADS",
        str(max(1, min(4, os.cpu_count() or 2)))
    )
)

os.environ["OMP_NUM_THREADS"] = str(CPU_THREADS)
os.environ["MKL_NUM_THREADS"] = str(CPU_THREADS)
os.environ["OPENBLAS_NUM_THREADS"] = str(CPU_THREADS)
os.environ["NUMEXPR_NUM_THREADS"] = str(CPU_THREADS)
os.environ["TOKENIZERS_PARALLELISM"] = "false"

import json
import math
import re
import threading
import time
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F
from flask import Flask, Response, jsonify, request, send_file
from huggingface_hub import HfApi, hf_hub_download, login
from tokenizers import Tokenizer


# ─────────────────────────────────────────────────────────────────────────────
# RUNTIME CONFIGURATION
# ─────────────────────────────────────────────────────────────────────────────

torch.set_num_threads(CPU_THREADS)
torch.set_num_interop_threads(1)

if hasattr(torch.backends, "mkldnn"):
    torch.backends.mkldnn.enabled = True

HF_TOKEN = os.environ.get("HF_TOKEN", "")
HF_DATA_REPO = os.environ.get(
    "DATA_REPO",
    "Bc-AI/MiniCoder-Instruction-dataset",
)
HF_MODEL_REPO = os.environ.get(
    "MODEL_REPO",
    "hugging-science/CodVa-2-DPO-ses4",
)

# INT8 significantly reduces model memory and can improve CPU inference.
# Disable with ENABLE_INT8=0 if it is slower on your particular CPU.
ENABLE_INT8 = os.environ.get("ENABLE_INT8", "1").lower() not in {
    "0",
    "false",
    "no",
    "off",
}

# Set to 2 or 4 to reduce SSE/frontend overhead. A value of 1 streams every token.
STREAM_EVERY = max(1, int(os.environ.get("STREAM_EVERY", "1")))

MAX_CTX = 2048
TARGET_SFT_STEPS = 20_000

if HF_TOKEN:
    login(token=HF_TOKEN)


# ─────────────────────────────────────────────────────────────────────────────
# MODEL CONFIGURATION
# ─────────────────────────────────────────────────────────────────────────────

@dataclass
class Config:
    vocab_size: int = 50304
    d_model: int = 896
    n_layers: int = 18
    n_heads: int = 14
    n_kv_heads: int = 2
    max_len: int = 2048
    rope_theta: float = 500_000.0
    window_size: int = 512
    pattern_mult: float = 2.5
    reason_mult: float = 0.75
    reason_depth: int = 2
    gate_hidden: int = 64
    gate_init: float = 0.0
    diff_lambda_init: float = 0.8
    tie_embeddings: bool = True

    @property
    def head_dim(self):
        return self.d_model // self.n_heads

    @property
    def pattern_dim(self):
        return (
            (int(self.d_model * self.pattern_mult) + 255) // 256
        ) * 256

    @property
    def reason_dim(self):
        return (
            (int(self.d_model * self.reason_mult) + 63) // 64
        ) * 64


# ─────────────────────────────────────────────────────────────────────────────
# MODEL COMPONENTS
# ─────────────────────────────────────────────────────────────────────────────

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.w = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        x32 = x.float()
        normed = x32 * torch.rsqrt(
            x32.pow(2).mean(dim=-1, keepdim=True) + self.eps
        )
        return (normed * self.w).to(x.dtype)


def precompute_rope(head_dim, max_len, theta, device):
    inv_freq = 1.0 / (
        theta ** (
            torch.arange(
                0,
                head_dim,
                2,
                device=device,
                dtype=torch.float32,
            ) / head_dim
        )
    )

    positions = torch.arange(
        max_len,
        device=device,
        dtype=torch.float32,
    )

    freqs = torch.outer(positions, inv_freq)
    return freqs.cos(), freqs.sin()


def apply_rope(x, cos, sin):
    """
    x:   [batch, heads, length, head_dim]
    cos: [length, head_dim / 2]
    sin: [length, head_dim / 2]
    """
    half = x.shape[-1] // 2

    cos = cos.unsqueeze(0).unsqueeze(0)
    sin = sin.unsqueeze(0).unsqueeze(0)

    x1 = x[..., :half]
    x2 = x[..., half:]

    return torch.cat(
        [
            x1 * cos - x2 * sin,
            x2 * cos + x1 * sin,
        ],
        dim=-1,
    )


def repeat_kv(x, n_rep):
    if n_rep == 1:
        return x

    batch, heads, length, head_dim = x.shape

    return (
        x.unsqueeze(2)
        .expand(batch, heads, n_rep, length, head_dim)
        .reshape(batch, heads * n_rep, length, head_dim)
    )


class DifferentialAttention(nn.Module):
    """
    Differential attention with preallocated KV caching.

    The cache stores unrepeated KV heads. This saves memory compared with
    storing the repeated GQA representation.
    """

    def __init__(self, cfg, layer_idx, local=False):
        super().__init__()

        self.n_pairs = cfg.n_heads // 2
        self.n_kv_pairs = max(1, cfg.n_kv_heads // 2)
        self.n_rep = self.n_pairs // self.n_kv_pairs

        self.head_dim = cfg.head_dim
        self.local = local
        self.window = cfg.window_size
        self.max_len = cfg.max_len

        d_model = cfg.d_model

        self.wq = nn.Linear(
            d_model,
            2 * self.n_pairs * self.head_dim,
            bias=False,
        )
        self.wk = nn.Linear(
            d_model,
            2 * self.n_kv_pairs * self.head_dim,
            bias=False,
        )
        self.wv = nn.Linear(
            d_model,
            self.n_kv_pairs * self.head_dim,
            bias=False,
        )
        self.wo = nn.Linear(
            self.n_pairs * self.head_dim,
            d_model,
            bias=False,
        )

        self.q_norm = RMSNorm(self.head_dim)
        self.k_norm = RMSNorm(self.head_dim)
        self.out_norm = RMSNorm(self.head_dim)

        self.lambda1 = nn.Parameter(torch.tensor(0.0))
        self.lambda2 = nn.Parameter(torch.tensor(0.0))

    def create_cache(self, batch_size, device, dtype):
        shape = (
            batch_size,
            self.n_kv_pairs,
            self.max_len,
            self.head_dim,
        )

        return {
            "k1": torch.empty(shape, device=device, dtype=dtype),
            "k2": torch.empty(shape, device=device, dtype=dtype),
            "v": torch.empty(shape, device=device, dtype=dtype),
            "length": 0,
        }

    def forward(self, x, cos, sin, cache, start_pos):
        batch, new_length, _ = x.shape
        head_dim = self.head_dim
        end_pos = start_pos + new_length

        if end_pos > self.max_len:
            raise ValueError(
                f"Context length {end_pos} exceeds maximum "
                f"context {self.max_len}"
            )

        q_all = self.wq(x).view(
            batch,
            new_length,
            self.n_pairs,
            2,
            head_dim,
        ).transpose(1, 2)

        k_all = self.wk(x).view(
            batch,
            new_length,
            self.n_kv_pairs,
            2,
            head_dim,
        ).transpose(1, 2)

        value = self.wv(x).view(
            batch,
            new_length,
            self.n_kv_pairs,
            head_dim,
        ).transpose(1, 2)

        q1 = q_all[..., 0, :]
        q2 = q_all[..., 1, :]
        k1 = k_all[..., 0, :]
        k2 = k_all[..., 1, :]

        q1 = self.q_norm(q1)
        q2 = self.q_norm(q2)
        k1 = self.k_norm(k1)
        k2 = self.k_norm(k2)

        q1 = apply_rope(q1, cos, sin)
        q2 = apply_rope(q2, cos, sin)
        k1 = apply_rope(k1, cos, sin)
        k2 = apply_rope(k2, cos, sin)

        # Write the new KV tensors into preallocated storage.
        cache["k1"][:, :, start_pos:end_pos, :].copy_(k1)
        cache["k2"][:, :, start_pos:end_pos, :].copy_(k2)
        cache["v"][:, :, start_pos:end_pos, :].copy_(value)
        cache["length"] = end_pos

        # During one-token decoding, local layers only need the most recent
        # window. During prompt prefill, retain the full prompt because every
        # prompt position needs its own local context.
        if self.local and new_length == 1:
            key_start = max(0, start_pos - self.window)
        else:
            key_start = 0

        key_end = end_pos

        cached_k1 = cache["k1"][:, :, key_start:key_end, :]
        cached_k2 = cache["k2"][:, :, key_start:key_end, :]
        cached_v = cache["v"][:, :, key_start:key_end, :]

        cached_k1 = repeat_kv(cached_k1, self.n_rep)
        cached_k2 = repeat_kv(cached_k2, self.n_rep)
        cached_v = repeat_kv(cached_v, self.n_rep)

        query_positions = torch.arange(
            start_pos,
            end_pos,
            device=x.device,
        ).unsqueeze(1)

        key_positions = torch.arange(
            key_start,
            key_end,
            device=x.device,
        ).unsqueeze(0)

        # For PyTorch SDPA boolean masks, True means the element participates.
        allowed = key_positions <= query_positions

        if self.local:
            allowed = allowed & (
                (query_positions - key_positions) <= self.window
            )

        out1 = F.scaled_dot_product_attention(
            q1,
            cached_k1,
            cached_v,
            attn_mask=allowed,
            dropout_p=0.0,
        )

        out2 = F.scaled_dot_product_attention(
            q2,
            cached_k2,
            cached_v,
            attn_mask=allowed,
            dropout_p=0.0,
        )

        differential_lambda = (
            torch.exp(self.lambda1)
            - torch.exp(self.lambda2)
            + 0.5
        )

        output = out1 - differential_lambda * out2

        output = (
            self.out_norm(output)
            .transpose(1, 2)
            .contiguous()
            .view(batch, new_length, -1)
        )

        return self.wo(output)


class DualPathFFN(nn.Module):
    def __init__(self, cfg):
        super().__init__()

        d_model = cfg.d_model
        pattern_dim = cfg.pattern_dim
        reason_dim = cfg.reason_dim

        self.pat_gate = nn.Linear(
            d_model,
            pattern_dim,
            bias=False,
        )
        self.pat_up = nn.Linear(
            d_model,
            pattern_dim,
            bias=False,
        )
        self.pat_down = nn.Linear(
            pattern_dim,
            d_model,
            bias=False,
        )

        reason_layers = [
            nn.Linear(d_model, reason_dim, bias=False),
            nn.SiLU(),
        ]

        for _ in range(cfg.reason_depth - 1):
            reason_layers.extend(
                [
                    nn.Linear(reason_dim, reason_dim, bias=False),
                    nn.SiLU(),
                ]
            )

        reason_layers.append(
            nn.Linear(reason_dim, d_model, bias=False)
        )

        self.reason = nn.Sequential(*reason_layers)
        self.merge = nn.Parameter(torch.zeros(d_model))

    def forward(self, x):
        pattern = self.pat_down(
            F.silu(self.pat_gate(x)) * self.pat_up(x)
        )

        weight = torch.sigmoid(self.merge)

        return (
            weight * pattern
            + (1.0 - weight) * self.reason(x)
        )


class TokenImportanceGate(nn.Module):
    def __init__(self, cfg):
        super().__init__()

        self.net = nn.Sequential(
            nn.Linear(
                cfg.d_model,
                cfg.gate_hidden,
                bias=True,
            ),
            nn.SiLU(),
            nn.Linear(
                cfg.gate_hidden,
                1,
                bias=True,
            ),
        )

    def forward(self, x):
        return x * (0.5 + torch.sigmoid(self.net(x)))


class Block(nn.Module):
    def __init__(self, cfg, layer_idx):
        super().__init__()

        self.norm1 = RMSNorm(cfg.d_model)
        self.norm2 = RMSNorm(cfg.d_model)

        self.attn = DifferentialAttention(
            cfg,
            layer_idx,
            local=(layer_idx % 2 == 0),
        )

        self.ffn = DualPathFFN(cfg)

    def create_cache(self, batch_size, device, dtype):
        return self.attn.create_cache(
            batch_size,
            device,
            dtype,
        )

    def forward(self, x, cos, sin, cache, start_pos):
        attention_output = self.attn(
            self.norm1(x),
            cos,
            sin,
            cache,
            start_pos,
        )

        x = x + attention_output
        x = x + self.ffn(self.norm2(x))

        return x


class CodVa2(nn.Module):
    def __init__(self, cfg):
        super().__init__()

        self.cfg = cfg

        self.embed = nn.Embedding(
            cfg.vocab_size,
            cfg.d_model,
        )

        self.importance = TokenImportanceGate(cfg)

        self.blocks = nn.ModuleList(
            [
                Block(cfg, layer_index)
                for layer_index in range(cfg.n_layers)
            ]
        )

        self.final_norm = RMSNorm(cfg.d_model)

        self.register_buffer(
            "rope_cos",
            torch.zeros(
                cfg.max_len,
                cfg.head_dim // 2,
            ),
        )

        self.register_buffer(
            "rope_sin",
            torch.zeros(
                cfg.max_len,
                cfg.head_dim // 2,
            ),
        )

        self._rope_ready = False

    def _init_rope(self, device):
        cos, sin = precompute_rope(
            self.cfg.head_dim,
            self.cfg.max_len,
            self.cfg.rope_theta,
            device,
        )

        self.rope_cos.copy_(cos)
        self.rope_sin.copy_(sin)
        self._rope_ready = True

    def create_cache(self, batch_size, device, dtype):
        return [
            block.create_cache(
                batch_size,
                device,
                dtype,
            )
            for block in self.blocks
        ]

    def forward(
        self,
        tokens,
        caches=None,
        start_pos=0,
        last_logits_only=True,
    ):
        batch, length = tokens.shape
        device = tokens.device

        if start_pos + length > self.cfg.max_len:
            raise ValueError(
                f"Requested position {start_pos + length}, but "
                f"max_len is {self.cfg.max_len}"
            )

        if not self._rope_ready:
            self._init_rope(device)

        cos = self.rope_cos[start_pos:start_pos + length]
        sin = self.rope_sin[start_pos:start_pos + length]

        x = self.importance(self.embed(tokens))

        if caches is None:
            caches = self.create_cache(
                batch_size=batch,
                device=x.device,
                dtype=x.dtype,
            )

        for layer_index, block in enumerate(self.blocks):
            x = block(
                x,
                cos,
                sin,
                caches[layer_index],
                start_pos,
            )

        x = self.final_norm(x)

        # Generation only needs the logits for the newest position. This avoids
        # a very expensive [prompt_length x vocabulary] projection on prefill.
        if last_logits_only:
            x = x[:, -1:, :]

        logits = F.linear(x, self.embed.weight)

        return logits, caches


# ─────────────────────────────────────────────────────────────────────────────
# TOKENIZER
# ─────────────────────────────────────────────────────────────────────────────

print("[init] loading tokenizer...")

tokenizer_path = hf_hub_download(
    repo_id=HF_DATA_REPO,
    filename="codva2_sft_tokenizer.json",
    repo_type="dataset",
    token=HF_TOKEN or None,
)

tokenizer = Tokenizer.from_file(tokenizer_path)

IM_START_ID = tokenizer.token_to_id("<|im_start|>")
IM_END_ID = tokenizer.token_to_id("<|im_end|>")

EOS_ID = tokenizer.token_to_id("<|endoftext|>")
if EOS_ID is None:
    EOS_ID = 50256

STOP_IDS = {
    token_id
    for token_id in (IM_END_ID, EOS_ID)
    if token_id is not None
}


# ─────────────────────────────────────────────────────────────────────────────
# CHECKPOINT HELPERS
# ─────────────────────────────────────────────────────────────────────────────

hf_api = HfApi()


def extract_checkpoint_step(filename):
    match = re.search(r"step[_-]?(\d+)", filename, flags=re.IGNORECASE)

    if match:
        return int(match.group(1))

    return 0


def get_latest_checkpoint():
    files = list(
        hf_api.list_repo_files(
            repo_id=HF_MODEL_REPO,
            repo_type="model",
            token=HF_TOKEN or None,
        )
    )

    checkpoints = [
        filename
        for filename in files
        if filename.endswith(".pt")
        and "config" not in filename
        and "final" not in filename
        and re.search(r"step[_-]?\d+", filename, re.IGNORECASE)
    ]

    if not checkpoints:
        return None, 0

    checkpoints.sort(key=extract_checkpoint_step)

    latest = checkpoints[-1]
    return latest, extract_checkpoint_step(latest)


# ─────────────────────────────────────────────────────────────────────────────
# LOAD MODEL
# ─────────────────────────────────────────────────────────────────────────────

print("[init] creating model...")

cfg = Config()
cfg.vocab_size = (
    (tokenizer.get_vocab_size() + 63) // 64
) * 64

float_model = CodVa2(cfg)

current_checkpoint = None
current_step = 0

try:
    current_checkpoint, current_step = get_latest_checkpoint()
except Exception as error:
    print(f"[init] checkpoint listing error: {error}")

if current_checkpoint:
    print(
        f"[init] loading checkpoint: "
        f"{current_checkpoint} (step {current_step})"
    )

    weights_path = hf_hub_download(
        repo_id=HF_MODEL_REPO,
        filename=current_checkpoint,
        repo_type="model",
        token=HF_TOKEN or None,
    )

    checkpoint = torch.load(
        weights_path,
        map_location="cpu",
        weights_only=False,
    )

    state_dict = (
        checkpoint["model"]
        if isinstance(checkpoint, dict) and "model" in checkpoint
        else checkpoint
    )

    float_model.load_state_dict(
        state_dict,
        strict=True,
    )

    del checkpoint
    del state_dict
else:
    print("[init] WARNING: no SFT checkpoint found")

float_model.eval()

if ENABLE_INT8:
    print("[init] applying dynamic INT8 quantization...")

    try:
        from torch.ao.quantization import quantize_dynamic

        model = quantize_dynamic(
            float_model,
            {nn.Linear},
            dtype=torch.qint8,
            inplace=False,
        )

        model.eval()
        del float_model

        quantization_name = "dynamic-int8"
        print("[init] dynamic INT8 enabled")
    except Exception as error:
        print(f"[init] INT8 quantization failed: {error}")
        print("[init] continuing with FP32")

        model = float_model
        quantization_name = "fp32"
else:
    model = float_model
    quantization_name = "fp32"

parameter_count = sum(
    parameter.numel()
    for parameter in model.parameters()
)

print(
    f"[init] ready | "
    f"{parameter_count / 1e6:.1f}M visible parameters | "
    f"step {current_step} | "
    f"{quantization_name} | "
    f"{CPU_THREADS} CPU threads | "
    f"KV cache enabled"
)


# ─────────────────────────────────────────────────────────────────────────────
# PROMPT AND SAMPLING
# ─────────────────────────────────────────────────────────────────────────────

SYSTEM_PROMPT = (
    "You are CodVa-2, a coding-specialized AI assistant. "
    "You write clean, correct, well-commented code. "
    "When solving problems, think step by step."
)


def format_prompt(user_message):
    return (
        f"<|im_start|>system\n"
        f"{SYSTEM_PROMPT}"
        f"<|im_end|>\n"
        f"<|im_start|>user\n"
        f"{user_message}"
        f"<|im_end|>\n"
        f"<|im_start|>assistant\n"
    )


def sample_token(logits, temperature, top_p, top_k):
    """
    Apply top-k first so top-p only sorts a small number of candidates.
    """

    if temperature <= 0.0:
        return int(torch.argmax(logits).item())

    logits = logits.float()
    logits = logits / max(temperature, 1e-5)

    vocabulary_size = logits.shape[-1]

    if top_k > 0:
        candidate_count = min(
            max(1, top_k),
            vocabulary_size,
        )

        candidate_logits, candidate_ids = torch.topk(
            logits,
            candidate_count,
        )
    else:
        candidate_logits = logits
        candidate_ids = torch.arange(
            vocabulary_size,
            device=logits.device,
        )

    probabilities = F.softmax(
        candidate_logits,
        dim=-1,
    )

    if 0.0 < top_p < 1.0:
        sorted_probabilities, sorted_positions = torch.sort(
            probabilities,
            descending=True,
        )

        cumulative_probabilities = torch.cumsum(
            sorted_probabilities,
            dim=-1,
        )

        remove = (
            cumulative_probabilities
            - sorted_probabilities
        ) > top_p

        sorted_probabilities = sorted_probabilities.masked_fill(
            remove,
            0.0,
        )

        probability_sum = sorted_probabilities.sum()

        if probability_sum <= 0:
            return int(
                candidate_ids[
                    torch.argmax(candidate_logits)
                ].item()
            )

        sorted_probabilities = (
            sorted_probabilities / probability_sum
        )

        sampled_sorted_position = torch.multinomial(
            sorted_probabilities,
            num_samples=1,
        )

        candidate_position = sorted_positions[
            sampled_sorted_position
        ]

        return int(
            candidate_ids[candidate_position].item()
        )

    sampled_position = torch.multinomial(
        probabilities,
        num_samples=1,
    )

    return int(
        candidate_ids[sampled_position].item()
    )


# ─────────────────────────────────────────────────────────────────────────────
# FLASK APPLICATION
# ─────────────────────────────────────────────────────────────────────────────

app = Flask(__name__)

stop_signal = threading.Event()

# CPU generation is serialized. Multiple simultaneous generations would divide
# CPU resources and make every request much slower.
generation_lock = threading.Lock()


@app.route("/")
def index():
    return send_file("index.html")


@app.route("/api/status", methods=["GET"])
def status():
    return jsonify(
        {
            "step": current_step,
            "progress": round(
                current_step / TARGET_SFT_STEPS * 100,
                1,
            ),
            "checkpoint": current_checkpoint or "none",
            "quantization": quantization_name,
            "cpu_threads": CPU_THREADS,
            "kv_cache": True,
            "stream_every": STREAM_EVERY,
        }
    )


@app.route("/stop", methods=["POST"])
def stop():
    stop_signal.set()
    return jsonify({"ok": True})


@app.route("/generate", methods=["POST"])
def generate():
    data = request.get_json(silent=True) or {}

    prompt = str(data.get("prompt") or "").strip()

    if not prompt:
        return jsonify({"error": "empty prompt"}), 400

    try:
        requested_max_new = int(data.get("max_new", 256))
        temperature = float(data.get("temperature", 0.3))
        top_p = float(data.get("top_p", 0.95))
        top_k = int(data.get("top_k", 50))
    except (TypeError, ValueError):
        return jsonify(
            {"error": "invalid generation settings"}
        ), 400

    requested_max_new = max(1, requested_max_new)
    temperature = max(0.0, temperature)
    top_p = min(max(top_p, 0.0), 1.0)
    top_k = max(0, top_k)

    stop_signal.clear()

    def stream():
        with generation_lock:
            full_prompt = format_prompt(prompt)
            prompt_ids = tokenizer.encode(full_prompt).ids

            # Always reserve at least one token for generation.
            prompt_ids = prompt_ids[-(MAX_CTX - 1):]

            available_new_tokens = MAX_CTX - len(prompt_ids)

            max_new = min(
                requested_max_new,
                available_new_tokens,
            )

            prompt_length = len(prompt_ids)
            generated_ids = []

            started_at = time.perf_counter()

            context_payload = {
                "type": "ctx",
                "prompt_tokens": prompt_length,
                "max_ctx": MAX_CTX,
                "max_new": max_new,
                "step": current_step,
                "quantization": quantization_name,
            }

            yield (
                f"data: {json.dumps(context_payload)}\n\n"
            )

            try:
                with torch.inference_mode():
                    # Prefill the complete prompt once.
                    input_tokens = torch.tensor(
                        [prompt_ids],
                        dtype=torch.long,
                    )

                    caches = None
                    start_pos = 0

                    for generation_index in range(max_new):
                        if stop_signal.is_set():
                            break

                        logits, caches = model(
                            input_tokens,
                            caches=caches,
                            start_pos=start_pos,
                            last_logits_only=True,
                        )

                        start_pos += input_tokens.shape[1]

                        next_token = sample_token(
                            logits[0, -1, :],
                            temperature=temperature,
                            top_p=top_p,
                            top_k=top_k,
                        )

                        generated_ids.append(next_token)

                        should_stream = (
                            len(generated_ids) == 1
                            or len(generated_ids) % STREAM_EVERY == 0
                            or next_token in STOP_IDS
                            or generation_index == max_new - 1
                        )

                        if should_stream:
                            raw_output = tokenizer.decode(
                                generated_ids,
                                skip_special_tokens=False,
                            )

                            elapsed = (
                                time.perf_counter()
                                - started_at
                            )

                            tokens_per_second = (
                                len(generated_ids)
                                / max(elapsed, 1e-6)
                            )

                            token_payload = {
                                "type": "token",
                                "raw": raw_output,
                                "tok_id": next_token,
                                "step": len(generated_ids),
                                "max_new": max_new,
                                "tps": round(
                                    tokens_per_second,
                                    2,
                                ),
                                "elapsed": round(
                                    elapsed,
                                    2,
                                ),
                                "total_tokens": (
                                    prompt_length
                                    + len(generated_ids)
                                ),
                                "prompt_tokens": prompt_length,
                                "max_ctx": MAX_CTX,
                                "model_step": current_step,
                            }

                            yield (
                                f"data: "
                                f"{json.dumps(token_payload)}"
                                f"\n\n"
                            )

                        if next_token in STOP_IDS:
                            break

                        if start_pos >= MAX_CTX:
                            break

                        # After prefill, only process the newest token.
                        input_tokens = torch.tensor(
                            [[next_token]],
                            dtype=torch.long,
                        )

                elapsed = time.perf_counter() - started_at

                tokens_per_second = (
                    len(generated_ids)
                    / max(elapsed, 1e-6)
                )

                done_payload = {
                    "type": "done",
                    "tps": round(tokens_per_second, 2),
                    "elapsed": round(elapsed, 2),
                    "tokens": len(generated_ids),
                    "stopped": stop_signal.is_set(),
                }

                yield (
                    f"data: {json.dumps(done_payload)}\n\n"
                )

            except GeneratorExit:
                stop_signal.set()
                return

            except Exception as error:
                print(
                    f"[generate] runtime error: "
                    f"{type(error).__name__}: {error}"
                )

                error_payload = {
                    "type": "error",
                    "error": str(error),
                }

                yield (
                    f"data: {json.dumps(error_payload)}\n\n"
                )

    return Response(
        stream(),
        mimetype="text/event-stream",
        headers={
            "Cache-Control": "no-cache, no-transform",
            "X-Accel-Buffering": "no",
            "Connection": "keep-alive",
        },
    )


if __name__ == "__main__":
    app.run(
        host="0.0.0.0",
        port=7860,
        threaded=True,
        use_reloader=False,
    )