#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ Inferno-Exp-1.7 — simple Gradio inference (loads like any Hub user, no pushing) ============================================================================== Loads Inferno-Exp-1.7 with plain Transformers (AutoModelForCausalLM + trust_remote_code=True) and serves a small Gradio app: open prompt box, sampling controls, streaming output with a Stop button, token/speed readout. Works with either source: • the Hub repo: MODEL_SOURCE = "AlgoDriveAI/Inferno-Exp-1.7" • the local export folder the push script builds (test before uploading): MODEL_SOURCE = "/content/Inferno-Exp-1.7-release" The public 1.7 model needs no token. For a private alternative, use an existing Hub login, set the HF_TOKEN environment variable, or set HF_TOKEN below. It is a base model: it continues text, it does not chat. Write the start of the text you want (e.g. a word problem followed by "Answer:"). Hardcoded constants, no argparse. Paste into one Colab cell and run; open the share link. """ # ============================================================================= # SETTINGS — edit these # ============================================================================= MODEL_SOURCE = "AlgoDriveAI/Inferno-Exp-1.7" # Hub repo id (or a local folder with the same files) MODEL_REVISION = "a24c67451dc7df3b3069e9ce54724ca815d96fda" # verified 1.7 weights HF_TOKEN = "" # optional; otherwise use env var / existing Hub login MODEL_DISPLAY_NAME = "Inferno-Exp-1.7" # Default slider values (all adjustable in the UI) DEF_MAX_NEW_TOKENS = 256 DEF_TEMPERATURE = 0.8 DEF_TOP_K = 40 DEF_TOP_P = 0.95 DEF_REPETITION_PENALTY = 1.1 DEF_SEED = -1 # -1 = different every run DTYPE = "auto" # "auto" (bf16 on newer GPUs, fp32 on CPU) / "float32" / "bfloat16" / "float16" GRADIO_SHARE = True # public *.gradio.live link (needed in Colab) GRADIO_PORT = 7860 INSTALL_DEPENDENCIES = True # ============================================================================= # RUN (no edits needed below) # ============================================================================= import os, sys, time, threading, subprocess, queue os.environ.setdefault("USE_TF", "0") os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") try: import google.colab # noqa: F401 IN_COLAB = True except ImportError: IN_COLAB = False if IN_COLAB and INSTALL_DEPENDENCIES: import importlib.metadata as _md def _version(name): try: return _md.version(name) except _md.PackageNotFoundError: return None _before = _version("transformers") subprocess.check_call([sys.executable, "-m", "pip", "install", "--quiet", "transformers==4.57.6", "safetensors", "huggingface_hub>=0.25,<1", "gradio>=4.44,<7", "sentencepiece", "protobuf"]) if _before != _version("transformers") and "transformers" in sys.modules: raise SystemExit("Transformers was just changed to a compatible version, but the old one is still loaded. " "Restart the session (Runtime → Restart session) and run this cell again.") import torch import gradio as gr from transformers import (AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer) DEVICE = "cuda" if torch.cuda.is_available() else "cpu" STOP_EVENT = threading.Event() def get_token(): # token=None lets the Hub use an existing login; public downloads work without one. return HF_TOKEN.strip() or os.environ.get("HF_TOKEN") or None def pick_dtype(): if DTYPE != "auto": return {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}[DTYPE] if DEVICE == "cuda": return torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16 return torch.float32 class StopOnEvent(StoppingCriteria): """Ends generate() as soon as the Stop button sets STOP_EVENT.""" def __call__(self, input_ids, scores, **kwargs): return STOP_EVENT.is_set() def load(): token = None if os.path.isdir(MODEL_SOURCE) else get_token() dtype = pick_dtype() revision = None if os.path.isdir(MODEL_SOURCE) else MODEL_REVISION tok = AutoTokenizer.from_pretrained(MODEL_SOURCE, revision=revision, token=token, use_fast=True) if tok.pad_token_id is None: tok.pad_token = tok.eos_token model = AutoModelForCausalLM.from_pretrained(MODEL_SOURCE, revision=revision, token=token, trust_remote_code=True, torch_dtype=dtype).to(DEVICE).eval() return model, tok, dtype MODEL, TOK, MODEL_DTYPE = load() CFG = MODEL.config N_PARAMS = sum(p.numel() for p in MODEL.parameters()) CONTEXT = CFG.max_position_embeddings print(f"Loaded {MODEL_SOURCE}: {N_PARAMS / 1e6:.1f}M params · {DEVICE} · {MODEL_DTYPE}") def generate(prompt, max_new, temperature, top_k, top_p, rep_penalty, seed): STOP_EVENT.clear() if not prompt: yield "", "Type a prompt first." return ids = TOK(prompt, return_tensors="pt", add_special_tokens=False).to(DEVICE) room = CONTEXT - ids["input_ids"].shape[1] if room <= 0: yield prompt, f"Prompt is longer than the model's {CONTEXT}-token context; shorten it." return seed = int(seed) if seed >= 0: torch.manual_seed(seed) greedy = temperature <= 0 sampling = {} if greedy else dict(temperature=float(temperature), top_k=int(top_k), top_p=float(top_p)) streamer = TextIteratorStreamer(TOK, skip_prompt=False, skip_special_tokens=True, timeout=1.0) result = {} def worker(): try: with torch.inference_mode(): result["ids"] = MODEL.generate( **ids, max_new_tokens=min(int(max_new), room), do_sample=not greedy, repetition_penalty=float(rep_penalty), eos_token_id=TOK.eos_token_id, pad_token_id=TOK.pad_token_id, use_cache=False, streamer=streamer, stopping_criteria=StoppingCriteriaList([StopOnEvent()]), **sampling) except Exception as exc: result["error"] = exc streamer.end() # unblock the UI if the generation worker fails started = time.time() thread = threading.Thread(target=worker, daemon=True) thread.start() text = "" try: while True: try: piece = next(streamer) except queue.Empty: if not thread.is_alive(): break continue except StopIteration: break text += piece yield text, f"generating… {time.time() - started:.1f}s" thread.join() finally: if thread.is_alive(): STOP_EVENT.set() if "error" in result: exc = result["error"] yield text or prompt, f"Generation failed: {type(exc).__name__}: {exc}" return new = result["ids"].shape[1] - ids["input_ids"].shape[1] if "ids" in result else 0 took = time.time() - started status = "stopped" if STOP_EVENT.is_set() else "done" yield text or prompt, (f"**{status}** · {new} tokens · {took:.1f}s · {new / max(took, 1e-6):.1f} tok/s · " f"{'greedy' if greedy else f'T={temperature} · top-k {int(top_k)} · top-p {top_p}'} · seed {seed}") def output_box(): kwargs = dict(label="Output (prompt + continuation)", lines=16, max_lines=30, interactive=False) for extra in (dict(show_copy_button=True), dict(buttons=["copy"]), {}): try: return gr.Textbox(**kwargs, **extra) except TypeError: continue def build_ui(): gr_major = int(gr.__version__.split(".")[0]) theme = gr.themes.Soft(primary_hue="indigo", secondary_hue="blue") blocks_kwargs = {} if gr_major >= 6 else dict(theme=theme) # Gradio 6 moved theme to launch() launch_kwargs = dict(theme=theme) if gr_major >= 6 else {} with gr.Blocks(title=MODEL_DISPLAY_NAME, **blocks_kwargs) as demo: gr.Markdown(f"# 🔥 {MODEL_DISPLAY_NAME}\n" f"{N_PARAMS / 1e6:.1f}M-parameter math-first base model · `{CFG.num_hidden_layers}L × {CFG.hidden_size}d` · " f"{CFG.num_attention_heads}Q/{CFG.num_key_value_heads}KV heads · context {CONTEXT} · " f"continues text (not a chat model)") with gr.Row(): with gr.Column(scale=3): prompt = gr.Textbox(label="Prompt", lines=6, placeholder="Write the beginning of the text to continue…") with gr.Row(): run = gr.Button("Generate", variant="primary") stop = gr.Button("Stop", variant="stop") output = output_box() status = gr.Markdown("") with gr.Column(scale=1): max_new = gr.Slider(16, CONTEXT, value=min(DEF_MAX_NEW_TOKENS, CONTEXT), step=16, label="Max new tokens") temperature = gr.Slider(0.0, 1.5, value=DEF_TEMPERATURE, step=0.05, label="Temperature (0 = greedy)") top_k = gr.Slider(0, 200, value=DEF_TOP_K, step=1, label="Top-k (0 = off)") top_p = gr.Slider(0.05, 1.0, value=DEF_TOP_P, step=0.01, label="Top-p") rep = gr.Slider(1.0, 1.5, value=DEF_REPETITION_PENALTY, step=0.01, label="Repetition penalty") seed = gr.Number(value=DEF_SEED, precision=0, label="Seed (-1 = random)") with gr.Accordion("Model info", open=False): gr.Markdown(f"- **Source:** `{MODEL_SOURCE}`\n- **Parameters:** `{N_PARAMS:,}`\n" f"- **Device / dtype:** `{DEVICE}` / `{str(MODEL_DTYPE).replace('torch.', '')}`\n" f"- **Vocab:** `{CFG.vocab_size:,}` · EOS `{TOK.eos_token}`") inputs = [prompt, max_new, temperature, top_k, top_p, rep, seed] events = dict(inputs=inputs, outputs=[output, status], concurrency_limit=1, concurrency_id="inferno-generation") run.click(generate, **events) prompt.submit(generate, **events) stop.click(lambda: STOP_EVENT.set(), queue=False) return demo, launch_kwargs if __name__ == "__main__": demo, launch_kwargs = build_ui() demo.queue().launch(share=GRADIO_SHARE, server_port=GRADIO_PORT, server_name="0.0.0.0", **launch_kwargs)