Inferno-Exp-1.7 / inference /gradio_demo.py
AlgoDriveAI's picture
Add collapsible Colab Gradio inference demo for Inferno-Exp-1.7
0122846 verified
Raw History Blame Contribute Delete
10.5 kB
#!/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)