Upload app.py with huggingface_hub

app.py

@@ -0,0 +1,455 @@
+#!/usr/bin/env python3
+"""
+Markov Chain Language Model — Interactive Demo
+OpenTransformers Ltd | Part of AGILLM Research
+
+Classical n-gram LM with Modified Kneser-Ney smoothing.
+GPU hash tables (sorted int64 + searchsorted) for parallel inference.
+Runs on CPU for HF Spaces compatibility.
+"""
+
+import os, sys, math, time, pickle, gc
+from pathlib import Path
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+import gradio as gr
+from huggingface_hub import hf_hub_download
+
+# ─── Force CPU for HF Spaces ───
+DEV = torch.device("cpu")
+
+# ─── Tokenizer ───
+TOKENIZER_ID = os.environ.get("TOKENIZER_ID", "gpt2")
+
+def _load_tokenizer():
+    from transformers import AutoTokenizer, logging as hf_log
+    hf_log.set_verbosity_error()
+    t = AutoTokenizer.from_pretrained(TOKENIZER_ID, use_fast=True)
+    if t.pad_token is None:
+        t.add_special_tokens({"pad_token": "<|pad|>"})
+    return t
+
+tok = _load_tokenizer()
+VOCAB = max(tok.get_vocab().values()) + 1
+EOS = tok.eos_token_id if tok.eos_token_id is not None else tok.sep_token_id
+
+# ─── FNV-1a Hashing ───
+FNV_OFFSET = 14695981039346656037
+FNV_PRIME  = 1099511628211
+MASK64     = (1 << 64) - 1
+INT64_MAX  = (1 << 63) - 1
+INT64_WRAP = 1 << 64
+FNV_OFFSET_S = FNV_OFFSET - INT64_WRAP
+
+def _hash_ngram_batch_gpu(contexts: torch.Tensor) -> torch.Tensor:
+    B, N = contexts.shape
+    h = torch.full((B,), FNV_OFFSET_S, dtype=torch.int64, device=contexts.device)
+    for i in range(N):
+        h = h ^ contexts[:, i]
+        h = h * FNV_PRIME
+    return h
+
+
+class GPUHashTable:
+    """Immutable hash table using sorted int64 keys + searchsorted."""
+
+    def __init__(self):
+        self.hashes: Optional[torch.Tensor] = None
+        self.counts: Optional[torch.Tensor] = None
+        self.continuation_counts: Optional[torch.Tensor] = None
+        self.total: int = 0
+        self.size: int = 0
+
+    def batch_lookup(self, hashes: torch.Tensor) -> torch.Tensor:
+        if self.size == 0:
+            return torch.zeros_like(hashes)
+        idx = torch.searchsorted(self.hashes, hashes).clamp(0, self.size - 1)
+        found = (self.hashes[idx] == hashes)
+        return torch.where(found, self.counts[idx], torch.zeros_like(hashes))
+
+    def batch_lookup_continuations(self, hashes: torch.Tensor) -> torch.Tensor:
+        if self.size == 0 or self.continuation_counts is None:
+            return torch.zeros_like(hashes)
+        idx = torch.searchsorted(self.hashes, hashes).clamp(0, self.size - 1)
+        found = (self.hashes[idx] == hashes)
+        return torch.where(found, self.continuation_counts[idx], torch.zeros_like(hashes))
+
+    def memory_bytes(self) -> int:
+        total = 0
+        for t in [self.hashes, self.counts, self.continuation_counts]:
+            if t is not None: total += t.nelement() * t.element_size()
+        return total
+
+
+def _hash_ngram_py(ngram: tuple) -> int:
+    h = FNV_OFFSET
+    for t in ngram:
+        h ^= (t & MASK64)
+        h = (h * FNV_PRIME) & MASK64
+    return h if h <= INT64_MAX else h - INT64_WRAP
+
+
+class MarkovLM:
+    """N-gram LM with Modified Kneser-Ney smoothing."""
+
+    def __init__(self, max_order: int = 5):
+        self.max_order = max_order
+        self.cpu_counts: List[Dict] = []
+        self.total_tokens = 0
+        self.tokens_trained = 0
+        self.gpu_ngram_tables: List[Optional[GPUHashTable]] = [None] * max_order
+        self.gpu_context_tables: List[Optional[GPUHashTable]] = [None] * max_order
+        self.frozen = False
+        self.discounts: List[Tuple[float, float, float]] = [(0.5, 1.0, 1.5)] * max_order
+        self.gpu_unigram_probs: Optional[torch.Tensor] = None
+
+    def freeze(self, device=DEV, prune_threshold: int = 1):
+        print(f"[freeze] Building hash tables on {device}...")
+        t0 = time.time()
+
+        for order in range(self.max_order):
+            d = self.cpu_counts[order]
+            if not d:
+                self.gpu_ngram_tables[order] = GPUHashTable()
+                self.gpu_context_tables[order] = GPUHashTable()
+                continue
+
+            # Build ngram table
+            entries = []
+            for ctx, nexts in d.items():
+                for next_tok, cnt in nexts.items():
+                    if cnt >= prune_threshold or order <= 1:
+                        key = ctx + (next_tok,)
+                        h = _hash_ngram_py(key)
+                        entries.append((h, cnt))
+
+            gt = GPUHashTable()
+            if entries:
+                entries.sort(key=lambda x: x[0])
+                gt.hashes = torch.tensor([e[0] for e in entries], dtype=torch.int64, device=device)
+                gt.counts = torch.tensor([e[1] for e in entries], dtype=torch.int64, device=device)
+                gt.total = sum(e[1] for e in entries)
+                gt.size = len(entries)
+            self.gpu_ngram_tables[order] = gt
+
+            # Build context table
+            ctx_entries = []
+            for ctx, nexts in d.items():
+                h = _hash_ngram_py(ctx)
+                total = sum(nexts.values())
+                n_unique = len(nexts)
+                ctx_entries.append((h, total, n_unique))
+
+            ct = GPUHashTable()
+            if ctx_entries:
+                ctx_entries.sort(key=lambda x: x[0])
+                ct.hashes = torch.tensor([e[0] for e in ctx_entries], dtype=torch.int64, device=device)
+                ct.counts = torch.tensor([e[1] for e in ctx_entries], dtype=torch.int64, device=device)
+                ct.continuation_counts = torch.tensor([e[2] for e in ctx_entries], dtype=torch.int64, device=device)
+                ct.total = sum(e[1] for e in ctx_entries)
+                ct.size = len(ctx_entries)
+            self.gpu_context_tables[order] = ct
+
+            n_ent = gt.size
+            mem = (gt.memory_bytes() + ct.memory_bytes()) / 1e6
+            print(f"  {order+1}-gram: {n_ent:,} entries | {mem:.1f} MB")
+
+        # Estimate KN discounts
+        self._estimate_discounts()
+
+        # Unigram probs
+        if self.cpu_counts[0]:
+            uni = self.cpu_counts[0].get((), {})
+            probs = torch.zeros(VOCAB, dtype=torch.float32, device=device)
+            total = sum(uni.values())
+            if total > 0:
+                for tok_id, cnt in uni.items():
+                    if 0 <= tok_id < VOCAB:
+                        probs[tok_id] = cnt / total
+            self.gpu_unigram_probs = probs
+
+        # Free CPU dicts
+        self.cpu_counts = []
+        gc.collect()
+
+        self.frozen = True
+        elapsed = time.time() - t0
+        gpu_mem = sum(
+            (gt.memory_bytes() if gt else 0) + (ct.memory_bytes() if ct else 0)
+            for gt, ct in zip(self.gpu_ngram_tables, self.gpu_context_tables)
+        )
+        print(f"[freeze] Done in {elapsed:.1f}s | {gpu_mem/1e6:.1f} MB total")
+
+    def _estimate_discounts(self):
+        for order in range(self.max_order):
+            gt = self.gpu_ngram_tables[order]
+            if gt is None or gt.size == 0:
+                continue
+            counts = gt.counts
+            n1 = (counts == 1).sum().item()
+            n2 = (counts == 2).sum().item()
+            n3 = (counts == 3).sum().item()
+            n4 = (counts == 4).sum().item()
+            if n1 == 0 or n2 == 0:
+                continue
+            Y = n1 / (n1 + 2 * n2)
+            D1 = max(0.01, min(1 - 2 * Y * (n2 / max(n1, 1)), 0.99))
+            D2 = max(D1, min(2 - 3 * Y * (n3 / max(n2, 1)), 1.99))
+            D3 = max(D2, min(3 - 4 * Y * (n4 / max(n3, 1)), 2.99))
+            self.discounts[order] = (D1, D2, D3)
+
+    @torch.no_grad()
+    def batch_log_probs(self, contexts: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        B = targets.shape[0]
+        device = targets.device
+        log_probs = torch.full((B,), math.log(1.0 / VOCAB), dtype=torch.float32, device=device)
+
+        for order in range(self.max_order):
+            gt = self.gpu_ngram_tables[order]
+            ct = self.gpu_context_tables[order]
+            if gt is None or ct is None or gt.size == 0:
+                continue
+
+            ctx_len = order
+            if ctx_len == 0:
+                if self.gpu_unigram_probs is not None:
+                    safe_t = targets.clamp(0, VOCAB - 1)
+                    uni_p = self.gpu_unigram_probs[safe_t]
+                    valid = uni_p > 0
+                    log_probs = torch.where(valid, torch.log(uni_p + 1e-30), log_probs)
+                continue
+
+            if ctx_len > contexts.shape[1]:
+                continue
+
+            ctx = contexts[:, -ctx_len:]
+            has_ctx = (ctx >= 0).all(dim=1)
+            if not has_ctx.any():
+                continue
+
+            full_ngram = torch.cat([ctx, targets.unsqueeze(1)], dim=1)
+            ngram_counts = gt.batch_lookup(_hash_ngram_batch_gpu(full_ngram)).float()
+            ctx_hashes = _hash_ngram_batch_gpu(ctx)
+            ctx_totals = ct.batch_lookup(ctx_hashes).float()
+            ctx_uniques = ct.batch_lookup_continuations(ctx_hashes).float()
+
+            D1, D2, D3 = self.discounts[order]
+            discount = torch.where(ngram_counts >= 3, D3,
+                       torch.where(ngram_counts >= 2, D2,
+                       torch.where(ngram_counts >= 1, D1, 0.0)))
+
+            numerator = (ngram_counts - discount).clamp(min=0)
+            denominator = ctx_totals.clamp(min=1)
+            gamma = (D3 * ctx_uniques) / denominator
+            gamma = gamma.clamp(0, 1)
+            p_lower = log_probs.exp()
+            p_combined = numerator / denominator + gamma * p_lower
+            valid = has_ctx & (ctx_totals > 0)
+            log_probs = torch.where(valid, torch.log(p_combined.clamp(min=1e-30)), log_probs)
+
+        return log_probs
+
+    @torch.no_grad()
+    def generate(self, prompt: str, max_new: int = 200, temperature: float = 0.8,
+                 top_k: int = 50, top_p: float = 0.9):
+        assert self.frozen, "Call freeze() first"
+        ctx_len = self.max_order - 1
+        ids = tok.encode(prompt)
+        n_cands = min(top_k * 10, VOCAB)
+
+        if self.gpu_unigram_probs is not None:
+            _, candidates = self.gpu_unigram_probs.topk(n_cands)
+        else:
+            candidates = torch.arange(n_cands, device=DEV)
+
+        for _ in range(max_new):
+            if len(ids) >= ctx_len:
+                ctx = ids[-ctx_len:]
+            else:
+                ctx = [-1] * (ctx_len - len(ids)) + ids
+
+            ctx_t = torch.tensor([ctx], dtype=torch.int64, device=DEV).expand(n_cands, ctx_len)
+            log_probs = self.batch_log_probs(ctx_t, candidates)
+
+            probs = (log_probs / max(temperature, 1e-8)).softmax(0)
+
+            if top_k > 0 and top_k < n_cands:
+                vals, idx = probs.topk(top_k)
+                mask = torch.zeros_like(probs)
+                mask.scatter_(0, idx, vals)
+                probs = mask
+
+            if top_p < 1.0:
+                sp, si = probs.sort(descending=True)
+                cum = sp.cumsum(0)
+                cutoff = cum > top_p
+                cutoff[0] = False
+                sp[cutoff] = 0
+                probs = torch.zeros_like(probs).scatter_(0, si, sp)
+
+            if probs.sum() == 0:
+                next_tok = candidates[0].item()
+            else:
+                probs = probs / probs.sum()
+                next_tok = candidates[probs.multinomial(1).item()].item()
+
+            ids.append(next_tok)
+            if next_tok == EOS:
+                break
+
+        return tok.decode(ids, skip_special_tokens=True)
+
+    @classmethod
+    def load(cls, path: str) -> 'MarkovLM':
+        p = Path(path)
+        for suffix in ['.cpu.pkl', '.pkl']:
+            candidate = p.with_suffix(suffix)
+            if candidate.exists():
+                p = candidate; break
+
+        print(f"[load] {p} ({p.stat().st_size / 1e6:.1f} MB)...")
+        with open(p, 'rb') as f:
+            data = pickle.load(f)
+
+        model = cls(max_order=data['max_order'])
+        model.total_tokens = data['total_tokens']
+        model.tokens_trained = data['tokens_trained']
+        model.discounts = data.get('discounts', model.discounts)
+
+        for order in range(model.max_order):
+            raw = data['cpu_counts'][order]
+            dd = defaultdict(lambda: defaultdict(int))
+            for ctx, nexts in raw.items():
+                dd[ctx] = defaultdict(int, nexts)
+            model.cpu_counts.append(dd)
+
+        total_entries = sum(
+            sum(len(v) for v in model.cpu_counts[o].values())
+            for o in range(model.max_order)
+        )
+        print(f"[load] {model.max_order}-gram | {model.tokens_trained:,} tokens | {total_entries:,} entries")
+        print("[load] Freezing to CPU...")
+        model.freeze(device=DEV)
+        return model
+
+
+# ─── Load model at startup ───
+print("Downloading model from HuggingFace...")
+model_path = hf_hub_download(
+    repo_id="OpenTransformer/markov-5gram-500m",
+    filename="markov_5gram.cpu.pkl",
+    cache_dir="/tmp/markov_cache"
+)
+print(f"Loading model from {model_path}...")
+MODEL = MarkovLM.load(model_path)
+print("Model ready!")
+
+
+# ─── Gradio Interface ───
+def generate_text(prompt, max_tokens, temperature, top_k, top_p):
+    if not prompt.strip():
+        return "Please enter a prompt."
+    t0 = time.time()
+    result = MODEL.generate(
+        prompt=prompt,
+        max_new=int(max_tokens),
+        temperature=float(temperature),
+        top_k=int(top_k),
+        top_p=float(top_p),
+    )
+    elapsed = time.time() - t0
+    gen_tokens = len(tok.encode(result)) - len(tok.encode(prompt))
+    stats = f"\n\n---\n*Generated {gen_tokens} tokens in {elapsed:.2f}s ({gen_tokens/max(elapsed,0.01):.0f} tok/s)*"
+    return result + stats
+
+
+def get_model_info():
+    total_entries = sum(
+        gt.size for gt in MODEL.gpu_ngram_tables if gt
+    )
+    mem = sum(
+        (gt.memory_bytes() if gt else 0) + (ct.memory_bytes() if ct else 0)
+        for gt, ct in zip(MODEL.gpu_ngram_tables, MODEL.gpu_context_tables)
+    ) / 1e6
+
+    info = f"""## Model Statistics
+- **Architecture**: {MODEL.max_order}-gram with Modified Kneser-Ney smoothing
+- **Tokens trained**: {MODEL.tokens_trained:,}
+- **Total n-gram entries**: {total_entries:,}
+- **Memory usage**: {mem:.1f} MB
+- **Tokenizer**: GPT-2 ({VOCAB:,} vocab)
+- **Inference**: CPU (searchsorted batch lookup)
+
+### How it works
+This is a classical n-gram language model — no neural network, no parameters to learn via gradient descent.
+Instead, it counts how often sequences of tokens appear in the training data and uses those counts
+to predict the next token. Kneser-Ney smoothing interpolates between different context lengths
+(unigram through {MODEL.max_order}-gram) so that even unseen contexts get reasonable predictions.
+
+The n-gram counts are stored in sorted hash tables and looked up via binary search (`torch.searchsorted`),
+making inference parallel and efficient even on CPU.
+
+### Per-order breakdown"""
+
+    for order in range(MODEL.max_order):
+        gt = MODEL.gpu_ngram_tables[order]
+        if gt and gt.size > 0:
+            D1, D2, D3 = MODEL.discounts[order]
+            info += f"\n- **{order+1}-gram**: {gt.size:,} entries (D1={D1:.3f}, D2={D2:.3f}, D3+={D3:.3f})"
+
+    info += f"\n\n*Trained by [OpenTransformers Ltd](https://huggingface.co/OpenTransformer). Part of AGILLM research.*"
+    return info
+
+
+with gr.Blocks(
+    title="Markov Chain LM — OpenTransformers",
+    theme=gr.themes.Base(primary_hue="blue", neutral_hue="slate"),
+) as demo:
+    gr.Markdown("""# Markov Chain Language Model
+### Classical N-gram LM with Modified Kneser-Ney Smoothing
+*No neural network — pure statistical language modelling. [OpenTransformers Ltd](https://huggingface.co/OpenTransformer)*
+    """)
+
+    with gr.Row():
+        with gr.Column(scale=3):
+            prompt = gr.Textbox(
+                label="Prompt",
+                placeholder="Enter text to continue...",
+                lines=3,
+                value="The meaning of life is"
+            )
+            output = gr.Markdown(label="Generated Text")
+            generate_btn = gr.Button("Generate", variant="primary", size="lg")
+
+        with gr.Column(scale=1):
+            max_tokens = gr.Slider(10, 500, value=200, step=10, label="Max tokens")
+            temperature = gr.Slider(0.1, 2.0, value=0.8, step=0.1, label="Temperature")
+            top_k = gr.Slider(1, 200, value=50, step=1, label="Top-K")
+            top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-P")
+
+    generate_btn.click(
+        fn=generate_text,
+        inputs=[prompt, max_tokens, temperature, top_k, top_p],
+        outputs=output,
+    )
+
+    with gr.Accordion("Model Information", open=False):
+        gr.Markdown(get_model_info())
+
+    gr.Examples(
+        examples=[
+            ["The meaning of life is"],
+            ["In the beginning, there was"],
+            ["The president of the United States"],
+            ["Machine learning is a field of"],
+            ["Once upon a time in a land far away"],
+            ["The quick brown fox"],
+        ],
+        inputs=prompt,
+    )
+
+if __name__ == "__main__":
+    demo.launch()