Add standalone HF inference script (no nanochat dependency)

run_standalone_inference.py

@@ -0,0 +1,261 @@
+#!/usr/bin/env python3
+"""
+Standalone HF inference helper for HimalayaGPT models.
+
+This script intentionally depends only on:
+- torch
+- transformers
+- huggingface_hub
+
+No nanochat repo internals are required at runtime.
+"""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+from typing import List
+
+
+DEFAULT_PROMPTS = [
+    "नेपालको राजधानी के हो?",
+    "दुई वाक्यमा हिमालको महत्व बताऊ।",
+    "Write a short paragraph about machine learning.",
+    "What is 17 * 19? Show quick mental math.",
+    "Write a Python function to compute Fibonacci numbers.",
+]
+
+
+SPECIAL_TOKENS = [
+    "<|bos|>",
+    "<|user_start|>",
+    "<|user_end|>",
+    "<|assistant_start|>",
+    "<|assistant_end|>",
+    "<|python_start|>",
+    "<|python_end|>",
+    "<|output_start|>",
+    "<|output_end|>",
+]
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description="Run robust HF inference for HimalayaGPT")
+    p.add_argument("--repo-id", default="himalaya-ai/himalayagpt-0.5b-it")
+    p.add_argument("--revision", default="main")
+    p.add_argument("--force-download", action="store_true", help="Force fresh snapshot download from HF")
+    p.add_argument("--prompt-style", choices=["auto", "chat", "plain"], default="auto")
+    p.add_argument("--dtype", choices=["auto", "float32", "bfloat16"], default="auto")
+    p.add_argument("--temperature", type=float, default=0.8)
+    p.add_argument("--top-k", type=int, default=50)
+    p.add_argument("--max-new-tokens", type=int, default=96)
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument("--device", choices=["auto", "cuda", "cpu"], default="auto")
+    p.add_argument("--prompts-file", default=None, help="Optional .txt file with one prompt per line")
+    return p.parse_args()
+
+
+def _special_id(tokenizer, token: str) -> int | None:
+    special_map = getattr(tokenizer, "_special_to_id", None)
+    if isinstance(special_map, dict) and token in special_map:
+        tid = int(special_map[token])
+        if tid >= 0:
+            return tid
+    tid = tokenizer.convert_tokens_to_ids(token)
+    if tid is None:
+        return None
+    if tokenizer.unk_token_id is not None and tid == tokenizer.unk_token_id and token != tokenizer.unk_token:
+        return None
+    return int(tid)
+
+
+def _strip_special(text: str) -> str:
+    out = text
+    for tok in SPECIAL_TOKENS:
+        out = out.replace(tok, "")
+    return out.strip()
+
+
+def _load_prompts(prompts_file: str | None) -> List[str]:
+    if not prompts_file:
+        return list(DEFAULT_PROMPTS)
+    p = Path(prompts_file)
+    lines = [line.strip() for line in p.read_text(encoding="utf-8").splitlines() if line.strip()]
+    if not lines:
+        raise ValueError(f"No prompts in {p}")
+    return lines
+
+
+def _choose_device(arg: str):
+    import torch
+
+    if arg == "cuda":
+        if not torch.cuda.is_available():
+            raise RuntimeError("CUDA requested but not available")
+        return torch.device("cuda")
+    if arg == "cpu":
+        return torch.device("cpu")
+    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def _choose_torch_dtype(dtype_arg: str, device):
+    import torch
+
+    if dtype_arg == "float32":
+        return torch.float32
+    if dtype_arg == "bfloat16":
+        return torch.bfloat16
+    # auto policy: prefer bf16 if supported, otherwise fp32 (avoid fp16 instability)
+    if device.type == "cuda" and torch.cuda.is_bf16_supported():
+        return torch.bfloat16
+    return torch.float32
+
+
+def _resolve_prompt_style(style_arg: str, tokenizer) -> str:
+    if style_arg != "auto":
+        return style_arg
+    needed = ["<|user_start|>", "<|user_end|>", "<|assistant_start|>"]
+    if all(_special_id(tokenizer, tok) is not None for tok in needed):
+        return "chat"
+    return "plain"
+
+
+def _build_prompt_ids(tokenizer, prompt: str, prompt_style: str, vocab_size: int) -> List[int]:
+    user_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"]
+    if prompt_style == "plain":
+        ids = user_ids
+    else:
+        bos = _special_id(tokenizer, "<|bos|>")
+        u_s = _special_id(tokenizer, "<|user_start|>")
+        u_e = _special_id(tokenizer, "<|user_end|>")
+        a_s = _special_id(tokenizer, "<|assistant_start|>")
+        if None in (bos, u_s, u_e, a_s):
+            raise RuntimeError("Chat style requested but special tokens are missing")
+        ids = [int(bos), int(u_s)] + user_ids + [int(u_e), int(a_s)]
+
+    # hard clamp safety against malformed token ids
+    return [min(max(int(t), 0), vocab_size - 1) for t in ids]
+
+
+def _top_k_filter(logits, top_k: int):
+    import torch
+
+    if top_k <= 0:
+        return logits
+    k = min(top_k, logits.size(-1))
+    v, _ = torch.topk(logits, k)
+    masked = logits.clone()
+    masked[masked < v[:, [-1]]] = -float("inf")
+    return masked
+
+
+def generate_compat(model, input_ids, max_new_tokens: int, temperature: float, top_k: int, seed: int):
+    import torch
+    import torch.nn.functional as F
+
+    ids = input_ids
+    rng = None
+    if temperature > 0:
+        rng = torch.Generator(device=ids.device)
+        rng.manual_seed(seed)
+
+    for _ in range(max_new_tokens):
+        attention_mask = torch.ones_like(ids)
+        logits = model(input_ids=ids, attention_mask=attention_mask, return_dict=True).logits[:, -1, :]
+        logits = _top_k_filter(logits, top_k)
+        if temperature > 0:
+            probs = F.softmax(logits / temperature, dim=-1)
+            next_ids = torch.multinomial(probs, num_samples=1, generator=rng)
+        else:
+            next_ids = torch.argmax(logits, dim=-1, keepdim=True)
+        ids = torch.cat((ids, next_ids), dim=1)
+    return ids
+
+
+def main() -> None:
+    args = parse_args()
+
+    import torch
+    import transformers
+    from huggingface_hub import snapshot_download
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    if transformers.__version__ == "4.57.0":
+        print(
+            "[warn] transformers==4.57.0 is yanked on PyPI due to packaging issues. "
+            "Prefer transformers>=4.57.1."
+        )
+
+    prompts = _load_prompts(args.prompts_file)
+    device = _choose_device(args.device)
+    torch_dtype = _choose_torch_dtype(args.dtype, device)
+
+    local = snapshot_download(
+        repo_id=args.repo_id,
+        repo_type="model",
+        revision=args.revision,
+        force_download=args.force_download,
+    )
+    sha = Path(local).name
+
+    tok = AutoTokenizer.from_pretrained(local, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        local,
+        trust_remote_code=True,
+        torch_dtype=torch_dtype,
+        device_map="auto" if device.type == "cuda" else None,
+    )
+    if device.type == "cpu":
+        model = model.to(device)
+    model.eval()
+
+    cfg = model.config
+    vocab_size = int(getattr(cfg, "padded_vocab_size", getattr(cfg, "vocab_size", len(tok))))
+    context_window = int(
+        min(
+            x
+            for x in [
+                getattr(cfg, "sequence_len", None),
+                getattr(cfg, "max_position_embeddings", None),
+                getattr(tok, "model_max_length", None),
+            ]
+            if isinstance(x, int) and x > 0
+        )
+    )
+
+    prompt_style = _resolve_prompt_style(args.prompt_style, tok)
+    max_prompt_tokens = max(1, context_window - args.max_new_tokens)
+
+    print(f"repo={args.repo_id} revision={args.revision} snapshot_sha={sha}")
+    print(f"device={device} dtype={torch_dtype} prompt_style={prompt_style}")
+    print(f"context_window={context_window} max_prompt_tokens={max_prompt_tokens}")
+
+    torch.manual_seed(args.seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(args.seed)
+
+    for i, prompt in enumerate(prompts, 1):
+        p_ids = _build_prompt_ids(tok, prompt, prompt_style, vocab_size)[:max_prompt_tokens]
+        input_ids = torch.tensor([p_ids], dtype=torch.long, device=next(model.parameters()).device)
+
+        with torch.no_grad():
+            out = generate_compat(
+                model=model,
+                input_ids=input_ids,
+                max_new_tokens=args.max_new_tokens,
+                temperature=args.temperature,
+                top_k=args.top_k,
+                seed=args.seed + i,
+            )
+
+        completion_ids = out[0, input_ids.shape[1] :]
+        completion = tok.decode(completion_ids, skip_special_tokens=False)
+        completion = _strip_special(completion)
+
+        print(f"\n--- Prompt {i} ---")
+        print("Prompt:", prompt)
+        print("Completion:", completion if completion else "<empty>")
+
+
+if __name__ == "__main__":
+    main()