Upload 3 files

Dockerfile

@@ -0,0 +1,41 @@
+# ── Mochiva HF Space — CPU inference server ──────────────────────────────────
+# Base: Python 3.11 slim (small image, fast startup on HF free tier)
+FROM python:3.11-slim
+
+# HF Spaces runs as user 1000 — set up a non-root user
+RUN useradd -m -u 1000 mochiva
+WORKDIR /app
+RUN chown mochiva /app
+
+# ── System dependencies ────────────────────────────────────────────────────
+# Only what we strictly need: no CUDA, no build tools for heavy packages
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        git \
+        curl \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
+
+# ── Python dependencies ────────────────────────────────────────────────────
+COPY requirements.txt .
+RUN pip install --no-cache-dir --upgrade pip \
+ && pip install --no-cache-dir -r requirements.txt
+
+# ── App code ───────────────────────────────────────────────────────────────
+COPY app.py .
+
+# ── HF Spaces metadata ────────────────────────────────────────────────────
+# Port 7860 is the standard HF Space port
+EXPOSE 7860
+
+# ── Run as non-root ────────────────────────────────────────────────────────
+USER mochiva
+
+# ── Startup ────────────────────────────────────────────────────────────────
+# --workers 1: model is loaded once in the main process; threading handles concurrency
+# --timeout-keep-alive 30: keep SSE connections alive
+CMD ["uvicorn", "app:app", \
+     "--host", "0.0.0.0", \
+     "--port", "7860", \
+     "--workers", "1", \
+     "--timeout-keep-alive", "30", \
+     "--log-level", "info"]

app.py

@@ -0,0 +1,545 @@
+"""
+hf_space/app.py
+──────────────────────────────────────────────────────────────────────────────
+Mochiva inference server — runs on HuggingFace Spaces (free CPU tier).
+
+Architecture
+  • PyTorch re-implementation of the Mochiva model (mirrors train model.py)
+    — loads weights from safetensors exported by export.py
+  • FastAPI + Server-Sent Events (SSE) for streaming token-by-token responses
+  • Model + tokeniser loaded from HF Hub at startup
+  • Thread-safe: uses a queue to stream tokens from the generation thread
+
+Endpoints
+  POST /generate           — streaming SSE generation
+  POST /generate_full      — non-streaming, returns full response JSON
+  GET  /health             — liveness probe
+  GET  /info               — model metadata
+
+Environment variables
+  MODEL_REPO   : HF repo id  (default: "my-username/Mochiva-model")
+  HF_TOKEN     : optional HF token for private repos
+
+SSE protocol (matching the frontend expectation)
+  data: {"token": "...", "done": false}\n\n
+  data: {"token": "", "done": true}\n\n
+"""
+
+from __future__ import annotations
+import os
+import json
+import math
+import time
+import threading
+import queue
+from typing import Iterator, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel, Field
+
+from huggingface_hub import hf_hub_download, snapshot_download
+from tokenizers import Tokenizer
+
+
+# ─── Config ───────────────────────────────────────────────────────────────────
+
+MODEL_REPO = os.environ.get("MODEL_REPO", "my-username/Mochiva-model")
+HF_TOKEN   = os.environ.get("HF_TOKEN", None)
+DEVICE     = "cpu"
+MAX_CTX    = int(os.environ.get("MAX_CTX", "4096"))
+
+
+# ─── PyTorch model (mirrors Flax model in mochiva_training/model.py) ─────────
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps   = eps
+        self.scale = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        rms = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).sqrt()
+        return (x.float() / rms).to(x.dtype) * self.scale
+
+
+def precompute_freqs_cis(
+    head_dim: int,
+    max_seq: int,
+    theta: float = 10_000.0,
+    scaling_factor: float = 1.0,
+) -> torch.Tensor:
+    half   = head_dim // 2
+    freqs  = 1.0 / (theta ** (torch.arange(0, half, dtype=torch.float32) / half))
+    freqs  = freqs / scaling_factor
+    t      = torch.arange(max_seq, dtype=torch.float32)
+    freqs  = torch.outer(t, freqs)               # (seq, half)
+    return torch.polar(torch.ones_like(freqs), freqs)  # complex64
+
+
+def apply_rope(
+    xq: torch.Tensor,     # (B, T, nh, hd)
+    xk: torch.Tensor,
+    freqs_cis: torch.Tensor,  # (T, hd//2)  complex
+) -> tuple[torch.Tensor, torch.Tensor]:
+    def rotate(x):
+        x_c = x.float().reshape(*x.shape[:-1], x.shape[-1] // 2, 2)
+        x_c = torch.view_as_complex(x_c)         # (..., half)
+        fc  = freqs_cis.unsqueeze(0).unsqueeze(2) # (1, T, 1, half)
+        out = torch.view_as_real(x_c * fc).reshape(*x.shape)
+        return out.to(x.dtype)
+    return rotate(xq), rotate(xk)
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, cfg: dict):
+        super().__init__()
+        self.nh  = cfg["num_attention_heads"]
+        self.hd  = cfg["head_dim"]
+        H        = cfg["hidden_size"]
+        self.q_proj = nn.Linear(H, self.nh * self.hd, bias=False)
+        self.k_proj = nn.Linear(H, self.nh * self.hd, bias=False)
+        self.v_proj = nn.Linear(H, self.nh * self.hd, bias=False)
+        self.o_proj = nn.Linear(self.nh * self.hd, H, bias=False)
+
+    def forward(
+        self,
+        x: torch.Tensor,             # (B, T, H)
+        freqs_cis: torch.Tensor,     # (T, hd//2)
+        mask: torch.Tensor,          # (1, 1, T, T) bool
+        kv_cache: Optional[dict] = None,
+    ) -> torch.Tensor:
+        B, T, _ = x.shape
+        nh, hd  = self.nh, self.hd
+
+        q = self.q_proj(x).view(B, T, nh, hd)
+        k = self.k_proj(x).view(B, T, nh, hd)
+        v = self.v_proj(x).view(B, T, nh, hd)
+
+        q, k = apply_rope(q, k, freqs_cis)
+
+        if kv_cache is not None:
+            # Append current k, v to cache
+            if "k" in kv_cache:
+                k = torch.cat([kv_cache["k"], k], dim=1)
+                v = torch.cat([kv_cache["v"], v], dim=1)
+            kv_cache["k"] = k
+            kv_cache["v"] = v
+
+        # (B, nh, T, hd)
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+
+        scale = 1.0 / math.sqrt(hd)
+        attn  = torch.einsum("bhqd,bhkd->bhqk", q, k) * scale
+
+        # Apply causal mask (only over current q positions)
+        Tq, Tk = attn.shape[-2], attn.shape[-1]
+        if mask is not None:
+            m = mask[..., :Tq, :Tk]
+            attn = attn.masked_fill(~m, float("-inf"))
+
+        attn = F.softmax(attn.float(), dim=-1).to(q.dtype)
+        out  = torch.einsum("bhqk,bhkd->bhqd", attn, v)
+        out  = out.transpose(1, 2).contiguous().view(B, Tq, nh * hd)
+        return self.o_proj(out)
+
+
+class SwiGLUMLP(nn.Module):
+    def __init__(self, cfg: dict):
+        super().__init__()
+        H, I = cfg["hidden_size"], cfg["intermediate_size"]
+        self.gate_proj = nn.Linear(H, I, bias=False)
+        self.up_proj   = nn.Linear(H, I, bias=False)
+        self.down_proj = nn.Linear(I, H, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class MochivaBlock(nn.Module):
+    def __init__(self, cfg: dict):
+        super().__init__()
+        eps = cfg.get("rms_norm_eps", 1e-6)
+        self.attn_norm = RMSNorm(cfg["hidden_size"], eps)
+        self.mlp_norm  = RMSNorm(cfg["hidden_size"], eps)
+        self.attn      = CausalSelfAttention(cfg)
+        self.mlp       = SwiGLUMLP(cfg)
+
+    def forward(self, x, freqs_cis, mask, kv_cache=None):
+        x = x + self.attn(self.attn_norm(x), freqs_cis, mask, kv_cache)
+        x = x + self.mlp(self.mlp_norm(x))
+        return x
+
+
+class MochivaForInference(nn.Module):
+    """
+    Causal LM for inference.
+    Weights loaded from safetensors (exported by export.py).
+    Uses KV-cache for efficient auto-regressive decoding.
+    """
+
+    def __init__(self, cfg: dict):
+        super().__init__()
+        self.cfg = cfg
+        V = cfg["vocab_size"]
+        H = cfg["hidden_size"]
+        L = cfg["num_hidden_layers"]
+
+        self.embed_tokens = nn.Embedding(V, H)
+        self.layers       = nn.ModuleList([MochivaBlock(cfg) for _ in range(L)])
+        self.norm         = RMSNorm(H, cfg.get("rms_norm_eps", 1e-6))
+        # LM head is tied to embeddings — no extra parameter
+
+        hd  = cfg["head_dim"]
+        ctx = cfg["max_position_embeddings"]
+        theta  = cfg.get("rope_theta", 10_000.0)
+        scale  = cfg.get("rope_scaling_factor", 1.0)
+        freqs  = precompute_freqs_cis(hd, ctx, theta, scale)
+        self.register_buffer("freqs_cis", freqs)  # (ctx, hd//2)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,        # (B, T)
+        kv_caches: Optional[list] = None,
+    ) -> torch.Tensor:                  # (B, T, V)
+        B, T  = input_ids.shape
+
+        # If we have a KV cache, the position offset is the cached length
+        offset = 0
+        if kv_caches and "k" in kv_caches[0]:
+            offset = kv_caches[0]["k"].shape[1]
+
+        x = self.embed_tokens(input_ids)   # (B, T, H)
+
+        # Causal mask over full sequence (offset + T)
+        full_len = offset + T
+        mask     = torch.tril(torch.ones(full_len, full_len, dtype=torch.bool,
+                                         device=x.device))
+        mask     = mask.unsqueeze(0).unsqueeze(0)      # (1,1,full,full)
+
+        freqs = self.freqs_cis[offset : offset + T]
+
+        for i, layer in enumerate(self.layers):
+            kvc = kv_caches[i] if kv_caches else None
+            x   = layer(x, freqs, mask, kvc)
+
+        x      = self.norm(x)
+        logits = x @ self.embed_tokens.weight.T        # (B, T, V)
+        return logits
+
+    @torch.inference_mode()
+    def generate_stream(
+        self,
+        input_ids: torch.Tensor,    # (1, prompt_len)
+        max_new_tokens: int = 256,
+        temperature: float  = 0.8,
+        top_p: float        = 0.9,
+        top_k: int          = 50,
+        repetition_penalty: float = 1.1,
+        eos_token_id: int   = 2,
+    ) -> Iterator[int]:
+        """
+        Yields token IDs one by one.
+        Uses KV-cache for O(1) per-step memory after prompt encoding.
+        """
+        kv_caches = [{} for _ in self.layers]
+
+        # encode prompt
+        logits = self(input_ids, kv_caches)           # (1, T, V)
+        next_token = _sample(
+            logits[:, -1, :], temperature, top_p, top_k,
+            input_ids, repetition_penalty
+        )
+        yield int(next_token)
+
+        generated = input_ids.tolist()[0] + [int(next_token)]
+        cur = next_token.unsqueeze(0)
+
+        for _ in range(max_new_tokens - 1):
+            logits = self(cur, kv_caches)              # (1, 1, V)
+            next_token = _sample(
+                logits[:, -1, :], temperature, top_p, top_k,
+                torch.tensor([generated]), repetition_penalty
+            )
+            tok_id = int(next_token)
+            if tok_id == eos_token_id:
+                break
+            generated.append(tok_id)
+            yield tok_id
+            cur = next_token.unsqueeze(0)
+
+
+# ─── Sampling ─────────────────────────────────────────────────────────────────
+
+def _sample(
+    logits: torch.Tensor,         # (1, V)
+    temperature: float,
+    top_p: float,
+    top_k: int,
+    context_ids: torch.Tensor,
+    repetition_penalty: float,
+) -> torch.Tensor:
+    logits = logits.float().squeeze(0)    # (V,)
+
+    # repetition penalty
+    if repetition_penalty != 1.0:
+        for tok in set(context_ids.tolist()):
+            if logits[tok] < 0:
+                logits[tok] *= repetition_penalty
+            else:
+                logits[tok] /= repetition_penalty
+
+    if temperature < 1e-4:
+        return logits.argmax(keepdim=True)
+
+    logits = logits / temperature
+
+    # top-k
+    if top_k > 0:
+        v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+        logits[logits < v[-1]] = float("-inf")
+
+    # top-p (nucleus)
+    if top_p < 1.0:
+        sorted_logits, sorted_idx = torch.sort(logits, descending=True)
+        cum_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+        sorted_remove = cum_probs - F.softmax(sorted_logits, dim=-1) > top_p
+        sorted_logits[sorted_remove] = float("-inf")
+        logits = torch.zeros_like(logits).scatter_(0, sorted_idx, sorted_logits)
+
+    probs = F.softmax(logits, dim=-1)
+    return torch.multinomial(probs, num_samples=1)
+
+
+# ─── Weight loading ───────────────────────────────────────────────────────────
+
+def _remap_key(key: str) -> str:
+    """
+    Map flattened safetensors key → PyTorch nn.Module attribute path.
+    E.g. "embed_tokens/embedding" → "embed_tokens.weight"
+         "layer_0/attn/q_proj/kernel" → "layers.0.attn.q_proj.weight"
+    """
+    key = key.replace("/", ".")
+    key = key.replace("embed_tokens.embedding", "embed_tokens.weight")
+    # layer_N → layers.N
+    import re
+    key = re.sub(r"layer_(\d+)\.", r"layers.\1.", key)
+    # Flax kernel → PyTorch weight
+    key = key.replace(".kernel", ".weight")
+    # norms: scale → scale (already matches RMSNorm)
+    return key
+
+
+def load_weights(model: MochivaForInference, weights_path: str):
+    try:
+        from safetensors.torch import load_file
+        flat = load_file(weights_path, device=DEVICE)
+    except Exception:
+        # fallback: numpy npz
+        import numpy as np
+        npz  = np.load(weights_path)
+        flat = {k: torch.from_numpy(v) for k, v in npz.items()}
+
+    state_dict   = model.state_dict()
+    mapped       = {}
+    unmatched_st = []
+
+    for raw_key, tensor in flat.items():
+        pt_key = _remap_key(raw_key)
+        if pt_key in state_dict:
+            # Transpose: Flax Dense kernels are (in, out), PyTorch Linear (out, in)
+            if "weight" in pt_key and pt_key not in ("embed_tokens.weight",) \
+               and len(tensor.shape) == 2:
+                tensor = tensor.T
+            mapped[pt_key] = tensor.to(state_dict[pt_key].dtype)
+        else:
+            unmatched_st.append(pt_key)
+
+    # Tie LM head (no separate parameter)
+    missing, unexpected = model.load_state_dict(mapped, strict=False)
+    if missing:
+        print(f"[model] Missing keys: {missing[:5]}")
+    if unexpected:
+        print(f"[model] Unexpected keys: {unexpected[:5]}")
+    print(f"[model] Loaded {len(mapped)} tensors")
+
+
+# ─── Startup: load model ─────────────────────────────────────────────────────
+
+print(f"[startup] Downloading {MODEL_REPO} from HF Hub …")
+t0 = time.time()
+
+model_dir = snapshot_download(
+    MODEL_REPO,
+    token=HF_TOKEN,
+    ignore_patterns=["*.msgpack", "flax_model*"],
+)
+
+with open(f"{model_dir}/config.json") as f:
+    hf_cfg = json.load(f)
+
+with open(f"{model_dir}/special_tokens.json") as f:
+    special = json.load(f)
+
+tokenizer = Tokenizer.from_file(f"{model_dir}/tokenizer.json")
+BOS_ID    = special["bos_id"]
+EOS_ID    = special["eos_id"]
+PAD_ID    = special["pad_id"]
+
+with open(f"{model_dir}/generation_config.json") as f:
+    gen_cfg = json.load(f)
+
+model = MochivaForInference(hf_cfg)
+model.eval()
+
+weights_file = f"{model_dir}/model.safetensors"
+if not os.path.exists(weights_file):
+    weights_file = f"{model_dir}/model_weights.npz"
+
+load_weights(model, weights_file)
+print(f"[startup] Model ready in {time.time()-t0:.1f}s  "
+      f"(params: {sum(p.numel() for p in model.parameters())/1e6:.1f}M)")
+
+
+# ─── FastAPI ──────────────────────────────────────────────────────────────────
+
+app = FastAPI(title="Mochiva Inference", version="1.0.0")
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+# ─── Request / Response schemas ───────────────────────────────────────────────
+
+class GenerateRequest(BaseModel):
+    prompt: str
+    max_new_tokens: int  = Field(default=256, ge=1,  le=1024)
+    temperature:    float = Field(default=0.8,  ge=0.01, le=2.0)
+    top_p:          float = Field(default=0.9,  ge=0.0,  le=1.0)
+    top_k:          int   = Field(default=50,   ge=0,    le=500)
+    repetition_penalty: float = Field(default=1.1, ge=1.0, le=3.0)
+    mochi_name:     str   = ""    # injected persona context
+
+
+# ─── Streaming SSE endpoint ──────────────────────────────────────────────────
+
+def _sse_event(token: str = "", done: bool = False) -> str:
+    payload = json.dumps({"token": token, "done": done})
+    return f"data: {payload}\n\n"
+
+
+def _generate_sse(req: GenerateRequest) -> Iterator[str]:
+    # Build prompt with persona context if provided
+    prompt = req.prompt
+    if req.mochi_name:
+        prompt = (
+            f"<mochi>You are {req.mochi_name}, a cute and playful virtual pet "
+            f"called a Mochi. You are friendly, energetic, and love the person "
+            f"who takes care of you.</mochi> {prompt}"
+        )
+
+    ids = [BOS_ID] + tokenizer.encode(prompt).ids
+    if len(ids) > MAX_CTX - req.max_new_tokens:
+        ids = ids[-(MAX_CTX - req.max_new_tokens):]
+
+    input_ids = torch.tensor([ids], dtype=torch.long)
+
+    tok_queue: queue.Queue[Optional[int]] = queue.Queue()
+
+    def _worker():
+        try:
+            for tok_id in model.generate_stream(
+                input_ids,
+                max_new_tokens     = req.max_new_tokens,
+                temperature        = req.temperature,
+                top_p              = req.top_p,
+                top_k              = req.top_k,
+                repetition_penalty = req.repetition_penalty,
+                eos_token_id       = EOS_ID,
+            ):
+                tok_queue.put(tok_id)
+        finally:
+            tok_queue.put(None)   # sentinel
+
+    t = threading.Thread(target=_worker, daemon=True)
+    t.start()
+
+    buf = []
+    while True:
+        tok_id = tok_queue.get()
+        if tok_id is None:
+            break
+        buf.append(tok_id)
+        # Decode incrementally (handles multi-byte UTF-8 via backtrack)
+        text = tokenizer.decode(buf)
+        if text.endswith("▁") or text.endswith("Ġ"):
+            # incomplete byte — accumulate
+            continue
+        yield _sse_event(token=text)
+        buf = []
+
+    if buf:
+        yield _sse_event(token=tokenizer.decode(buf))
+    yield _sse_event(done=True)
+
+
+@app.post("/generate")
+def generate_stream(req: GenerateRequest):
+    return StreamingResponse(
+        _generate_sse(req),
+        media_type="text/event-stream",
+        headers={
+            "Cache-Control": "no-cache",
+            "X-Accel-Buffering": "no",
+        },
+    )
+
+
+# ─── Non-streaming endpoint ───────────────────────────────────────────────────
+
+@app.post("/generate_full")
+def generate_full(req: GenerateRequest):
+    tokens = []
+    for chunk in _generate_sse(req):
+        if chunk.startswith("data: "):
+            obj = json.loads(chunk[6:])
+            if not obj["done"]:
+                tokens.append(obj["token"])
+    return {"text": "".join(tokens), "model": MODEL_REPO}
+
+
+# ─── Health / info ────────────────────────────────────────────────────────────
+
+@app.get("/health")
+def health():
+    return {"status": "ok", "model": MODEL_REPO}
+
+
+@app.get("/info")
+def info():
+    return {
+        "model":      MODEL_REPO,
+        "vocab_size": hf_cfg["vocab_size"],
+        "layers":     hf_cfg["num_hidden_layers"],
+        "hidden":     hf_cfg["hidden_size"],
+        "context":    hf_cfg["max_position_embeddings"],
+        "device":     DEVICE,
+    }
+
+
+# ─── Entrypoint ───────────────────────────────────────────────────────────────
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)

requirements.txt

@@ -0,0 +1,17 @@
+# ── Mochiva HF Space — inference requirements ────────────────────────────────
+# CPU-only PyTorch (much smaller image than CUDA build)
+torch==2.3.0+cpu --extra-index-url https://download.pytorch.org/whl/cpu
+
+# Web server
+fastapi==0.111.0
+uvicorn[standard]==0.30.1
+pydantic==2.7.1
+
+# HF Hub for downloading the model at startup
+huggingface_hub==0.23.2
+
+# Fast BPE tokeniser (same library used at training time)
+tokenizers==0.19.1
+
+# Weights format
+safetensors==0.4.3