Spaces:

PrometheanStudio
/

Playground

Sleeping

Fixed app.py

fe1009e verified 23 days ago

19.1 kB

import json
import re
import shutil
import subprocess
import sys
from dataclasses import fields
from pathlib import Path

import gradio as gr
import spaces
import torch
from huggingface_hub import HfApi, hf_hub_download
from safetensors.torch import load_file


TITLE = "Promethean Studios Model Playground"
MODELS = {
    "Talos 10M": "PrometheanStudio/talos-10m",
    "Styx 100M": "PrometheanStudio/styx-100m",
}


TALOS_COMMIT = "e230fed9c18a5467e69ed0f15778a5b06d59cb5a"
TALOS_URL = "https://github.com/Promethean-Studios/talos.git"
CACHE = Path("/tmp/promethean-playground")
SOURCE = CACHE / "talos"
CACHE.mkdir(parents=True, exist_ok=True)

MODEL_CACHE = {}
TOKENIZER_CACHE = {}


FALLBACK_CONFIGS = {
    "PrometheanStudio/talos-10m": {
        "vocab_size": 1024,
        "max_seq_len": 512,
        "hidden_size": 448,
        "num_layers": 3,
        "num_attention_heads": 28,
        "num_kv_heads": 14,
        "head_dim": 16,
        "ffn_type": "dense",
        "intermediate_size": 1792,
        "attention_type": "full",
        "rope_theta": 10000.0,
        "layer_norm_eps": 1e-5,
        "tie_word_embeddings": False,
    },
    "PrometheanStudio/styx-100m": {
        "vocab_size": 1024,
        "max_seq_len": 512,
        "hidden_size": 1024,
        "num_layers": 6,
        "num_attention_heads": 64,
        "num_kv_heads": 32,
        "head_dim": 16,
        "ffn_type": "dense",
        "intermediate_size": 4096,
        "attention_type": "full",
        "rope_theta": 10000.0,
        "layer_norm_eps": 1e-5,
        "tie_word_embeddings": False,
    },
}


def repo_files(repo_id):
    return HfApi().list_repo_files(repo_id=repo_id, repo_type="model")


def find_tokenizer_file(files):
    candidates = [name for name in files if Path(name).name.lower() == "tokenizer.json"]
    if not candidates:
        raise FileNotFoundError("No tokenizer.json was found in the model repository.")
    candidates.sort(
        key=lambda name: (
            0 if name == "tokenizer/tokenizer.json" else 1,
            0 if name == "tokenizer.json" else 1,
            len(name),
            name,
        )
    )
    return candidates[0]


def ensure_talos_source():

    marker = SOURCE / ".talos_commit"
    required_files = (
        "model/gpt.py",
        "model/config.py",
        "inference/generate.py",
        "tokenizer/tokenizer.py",
        "tokenizer/vocab.py",
    )
    source_is_ready = (
        marker.is_file()
        and marker.read_text(encoding="utf-8").strip() == TALOS_COMMIT
        and all((SOURCE / filename).is_file() for filename in required_files)
    )

    if not source_is_ready:
        if SOURCE.exists():
            shutil.rmtree(SOURCE)
        SOURCE.mkdir(parents=True, exist_ok=True)
        commands = [
            ["git", "init", str(SOURCE)],
            ["git", "-C", str(SOURCE), "remote", "add", "origin", TALOS_URL],
            ["git", "-C", str(SOURCE), "fetch", "--depth", "1", "origin", TALOS_COMMIT],
            ["git", "-C", str(SOURCE), "checkout", "--detach", "FETCH_HEAD"],
        ]
        for command in commands:
            try:
                result = subprocess.run(
                    command, check=False, capture_output=True, text=True, timeout=240
                )
            except (OSError, subprocess.TimeoutExpired) as exc:
                raise RuntimeError(f"Could not prepare Talos source: {exc}") from exc
            if result.returncode != 0:
                detail = (result.stderr or result.stdout or "No git output").strip()
                raise RuntimeError(f"Talos git checkout failed: {detail[-1200:]}")
        if not all((SOURCE / filename).is_file() for filename in required_files):
            raise RuntimeError("The pinned Talos revision is missing expected source files.")
        marker.write_text(TALOS_COMMIT, encoding="utf-8")
        print(f"Talos source pinned to {TALOS_COMMIT}")

    source_prefix = str(SOURCE.resolve())
    sys.path[:] = [
        entry for entry in sys.path
        if str(Path(entry or ".").resolve()) != source_prefix
    ]
    sys.path.insert(0, source_prefix)


    for name, module in list(sys.modules.items()):
        if not (
            name in ("model", "inference", "tokenizer")
            or name.startswith(("model.", "inference.", "tokenizer."))
        ):
            continue
        module_file = getattr(module, "__file__", None)
        try:
            belongs_to_source = bool(
                module_file and Path(module_file).resolve().is_relative_to(SOURCE.resolve())
            )
        except (OSError, ValueError):
            belongs_to_source = False
        if not belongs_to_source:
            del sys.modules[name]


def load_tokenizer(repo_id, files):

    if repo_id in TOKENIZER_CACHE:
        return TOKENIZER_CACHE[repo_id]

    ensure_talos_source()
    from tokenizer.tokenizer import ByteLevelBPETokenizer

    filename = find_tokenizer_file(files)
    downloaded = hf_hub_download(repo_id=repo_id, filename=filename, repo_type="model")

    with open(downloaded, "r", encoding="utf-8") as handle:
        payload = json.load(handle)
    if not isinstance(payload, dict) or payload.get("type") != "ByteLevelBPE":
        raise RuntimeError(
            f"{repo_id}/{filename} is not the native Talos ByteLevelBPE format. "
            "A different tokenizer loader is required."
        )
    if payload.get("version") != 1:
        raise RuntimeError(
            f"Unsupported Talos tokenizer version {payload.get('version')!r} "
            f"in {repo_id}/{filename}."
        )

    tokenizer = ByteLevelBPETokenizer.from_file(str(downloaded))

    print(
        f"Tokenizer ready: {repo_id}/{filename}; "
        f"defined tokens={tokenizer.vocab_size}, "
        f"configured vocabulary={tokenizer.config.vocab_size}, "
        f"merges={tokenizer.merge_count}"
    )
    TOKENIZER_CACHE[repo_id] = tokenizer
    return tokenizer


def config_dict_from(value):
    if isinstance(value, dict):
        for nested_key in ("model_config", "config", "model"):
            nested = value.get(nested_key)
            if isinstance(nested, dict) and any(
                key in nested for key in ("hidden_size", "num_layers", "vocab_size")
            ):
                return nested
        return value
    if hasattr(value, "__dict__"):
        return vars(value)
    return None


def config_candidates(repo_id, files, checkpoint):
    from model.config import ModelConfig

    allowed = {item.name for item in fields(ModelConfig)}
    required = {
        "vocab_size", "max_seq_len", "hidden_size", "num_layers",
        "num_attention_heads", "num_kv_heads",
    }
    candidates = []

    if "config.json" in files:
        path = hf_hub_download(repo_id=repo_id, filename="config.json", repo_type="model")
        try:
            with open(path, "r", encoding="utf-8") as handle:
                raw = json.load(handle)
            normalized = config_dict_from(raw)
            if normalized:
                clean = {key: value for key, value in normalized.items() if key in allowed}
                if required.issubset(clean):
                    candidates.append(("repository config.json", clean))
        except Exception as exc:
            print(f"Could not use config.json for {repo_id}: {exc}")

    if isinstance(checkpoint, dict):
        for key in ("config", "model_config"):
            raw = config_dict_from(checkpoint.get(key))
            if raw:
                clean = {name: value for name, value in raw.items() if name in allowed}
                if required.issubset(clean):
                    candidates.append((f"checkpoint {key}", clean))

    if repo_id in FALLBACK_CONFIGS:
        candidates.append(("known model preset", FALLBACK_CONFIGS[repo_id]))

    unique, seen = [], set()
    for source, cfg in candidates:
        signature = json.dumps(cfg, sort_keys=True, default=str)
        if signature not in seen:
            unique.append((source, cfg))
            seen.add(signature)
    return unique


def read_weights(repo_id, files):
    safetensors_files = [name for name in files if name.endswith(".safetensors")]
    if safetensors_files:
        safetensors_files.sort(
            key=lambda name: (
                0 if name == "model.safetensors" else 1,
                len(name),
                name,
            )
        )
        filename = safetensors_files[0]
        path = hf_hub_download(repo_id=repo_id, filename=filename, repo_type="model")
        print(f"Loading inference weights: {repo_id}/{filename}")
        return load_file(path, device="cpu"), None

    checkpoint_files = [name for name in files if name.endswith((".pt", ".pth", ".bin"))]
    if not checkpoint_files:
        raise FileNotFoundError(f"No model weights found in {repo_id}.")

    def priority(name):
        basename = Path(name).name.lower()
        match = re.search(r"step[-_](\d+)", basename)
        step = int(match.group(1)) if match else -1
        return (
            0 if step >= 0 else 1,
            -step,
            0 if any(word in basename for word in ("final", "best", "model")) else 1,
            len(name),
            name,
        )

    checkpoint_files.sort(key=priority)
    filename = checkpoint_files[0]
    path = hf_hub_download(repo_id=repo_id, filename=filename, repo_type="model")
    print(f"Loading PyTorch checkpoint: {repo_id}/{filename}")

    try:
        checkpoint = torch.load(path, map_location="cpu", weights_only=True)
    except Exception as safe_exc:


        if repo_id not in MODELS.values():
            raise RuntimeError(
                f"Safe checkpoint loading failed: {safe_exc}. "
                "Publish a model.safetensors file instead."
            ) from safe_exc
        print(f"Safe loading failed for {filename}; trying legacy checkpoint reader")
        try:
            checkpoint = torch.load(path, map_location="cpu", weights_only=False)
        except Exception as legacy_exc:
            raise RuntimeError(
                f"Could not load {filename}: safe={safe_exc}; legacy={legacy_exc}"
            ) from legacy_exc

    state = checkpoint
    if isinstance(checkpoint, dict):
        for key in ("model", "model_state_dict", "state_dict", "model_state", "weights"):
            value = checkpoint.get(key)
            if isinstance(value, dict) and value and all(torch.is_tensor(v) for v in value.values()):
                state = value
                break
    if not isinstance(state, dict) or not state or not all(torch.is_tensor(v) for v in state.values()):
        raise RuntimeError(f"{filename} did not contain a model state dictionary.")
    return state, checkpoint if isinstance(checkpoint, dict) else None


def normalize_state_dict(state):
    for prefix in ("module.", "_orig_mod."):
        if state and all(key.startswith(prefix) for key in state):
            state = {key[len(prefix):]: value for key, value in state.items()}
    return state


def allowed_token_ids(tokenizer, model_vocab_size):


    text_ids = set(tokenizer.vocab.id_to_bytes.keys())
    valid_ids = text_ids | {int(tokenizer.eos_id)}
    if not valid_ids or min(valid_ids) < 0 or max(valid_ids) >= model_vocab_size:
        raise RuntimeError(
            "Tokenizer contains IDs outside the model vocabulary: "
            f"largest={max(valid_ids, default=-1)}, model_vocab={model_vocab_size}"
        )


    all_ids = text_ids | set(tokenizer.special_id_map().values())
    if min(all_ids, default=0) < 0 or max(all_ids, default=-1) >= model_vocab_size:
        raise RuntimeError(
            "Tokenizer special-token IDs exceed the model vocabulary: "
            f"largest={max(all_ids)}, model_vocab={model_vocab_size}"
        )
    return frozenset(valid_ids)


def load_model(repo_id):
    if repo_id in MODEL_CACHE:
        return MODEL_CACHE[repo_id]

    files = repo_files(repo_id)
    ensure_talos_source()
    tokenizer = load_tokenizer(repo_id, files)

    state, checkpoint = read_weights(repo_id, files)
    state = normalize_state_dict(state)

    from model import TalosGPT
    from model.config import ModelConfig

    attempts = []
    for source, raw_config in config_candidates(repo_id, files, checkpoint):
        try:
            config = ModelConfig(**raw_config).derive()
            model = TalosGPT(config)
            model.load_state_dict(state, strict=True)
            valid_ids = allowed_token_ids(tokenizer, config.vocab_size)
            model.eval()
            print(
                f"Loaded {repo_id} from {source}: "
                f"{model.num_parameters():,} parameters, "
                f"context={config.max_seq_len}, generatable_ids={len(valid_ids)}"
            )
            MODEL_CACHE[repo_id] = (model, tokenizer, valid_ids)
            return MODEL_CACHE[repo_id]
        except Exception as exc:
            attempts.append(
                f"{source} (hidden={raw_config.get('hidden_size')}, "
                f"layers={raw_config.get('num_layers')}): "
                f"{type(exc).__name__}: {str(exc)[:700]}"
            )

    raise RuntimeError(
        f"The weights in {repo_id} did not strictly match any Talos config.\n"
        + "\n".join(attempts)
    )


def get_model(label):
    if label not in MODELS:
        raise ValueError(f"Unknown model selection: {label!r}")
    return load_model(MODELS[label])


def choose_next_token(scores, history, temperature, top_p, repetition_penalty, allowed_mask):
    scores = scores.float().clone()

    if repetition_penalty > 1.0:
        for token_id in set(history[-128:]):
            if 0 <= token_id < scores.numel():
                if scores[token_id] < 0:
                    scores[token_id] *= repetition_penalty
                else:
                    scores[token_id] /= repetition_penalty

    scores = scores.masked_fill(~allowed_mask, float("-inf"))
    scores = torch.where(torch.isnan(scores), -torch.inf, scores)
    if not torch.isfinite(scores).any():
        raise RuntimeError("The model produced no finite logits for decodable tokens.")

    if temperature <= 0:
        return scores.argmax(dim=-1).view(1, 1)

    probabilities = torch.softmax(scores / max(float(temperature), 1e-5), dim=-1)
    if not torch.isfinite(probabilities).all():
        raise RuntimeError("Model output probabilities contain NaN or infinity.")

    sorted_probs, sorted_ids = torch.sort(probabilities, descending=True)
    cumulative = torch.cumsum(sorted_probs, dim=-1)
    remove = cumulative > float(top_p)
    remove[1:] = remove[:-1].clone()
    remove[0] = False
    sorted_probs[remove] = 0
    total = sorted_probs.sum()
    if not torch.isfinite(total) or total <= 0:
        return scores.argmax(dim=-1).view(1, 1)
    sampled_index = torch.multinomial(sorted_probs / total, num_samples=1)
    return sorted_ids.gather(0, sampled_index).view(1, 1)


@spaces.GPU
def generate(label, prompt, max_new_tokens, temperature, top_p, repetition_penalty):
    if not prompt or not prompt.strip():
        raise gr.Error("Enter a prompt first.")

    try:
        model, tokenizer, valid_ids = get_model(label)
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model = model.to(device).eval()


        prompt_ids = tokenizer.encode(prompt)
        if not prompt_ids:
            raise gr.Error("The tokenizer returned no tokens for the prompt.")

        max_new_tokens = int(max_new_tokens)
        max_seq_len = int(model.config.max_seq_len)
        if not 1 <= max_new_tokens < max_seq_len:
            raise gr.Error(f"New-token count must be between 1 and {max_seq_len - 1}.")
        prompt_ids = prompt_ids[-(max_seq_len - max_new_tokens):]
        if min(prompt_ids) < 0 or max(prompt_ids) >= model.config.vocab_size:
            raise RuntimeError("Encoded prompt contains token IDs outside the model vocabulary.")

        input_ids = torch.tensor([prompt_ids], dtype=torch.long, device=device)
        allowed_mask = torch.zeros(model.config.vocab_size, dtype=torch.bool, device=device)
        allowed_mask[list(valid_ids)] = True
        history = list(prompt_ids)
        generated_ids = []

        from inference.generate import prefill, decode_step

        with torch.inference_mode():
            logits, cache = prefill(model, input_ids)
            for step in range(max_new_tokens):
                scores = logits[0, -1, :]
                next_token = choose_next_token(
                    scores,
                    history,
                    float(temperature),
                    float(top_p),
                    float(repetition_penalty),
                    allowed_mask,
                )
                token_id = int(next_token.item())
                if token_id == tokenizer.eos_id:
                    break
                generated_ids.append(token_id)
                history.append(token_id)
                if step + 1 < max_new_tokens:

                    logits = decode_step(model, cache, next_token, position=cache.length)

        return tokenizer.decode(generated_ids)

    except gr.Error:
        raise
    except Exception as exc:
        print(f"Generation failure for {label}: {type(exc).__name__}: {exc}")
        raise gr.Error(f"Could not load or run {label}: {type(exc).__name__}: {exc}") from exc


with gr.Blocks(title=TITLE) as demo:
    gr.Markdown(
        "# Promethean Studios Model Playground\n"
        "Experiment with Talos and Styx checkpoints. These are research base models, "
        "not guaranteed to behave like instruction-tuned chat assistants."
    )

    with gr.Row():
        model_choice = gr.Dropdown(
            choices=list(MODELS.keys()), value="Styx 100M", label="Model"
        )
        max_tokens = gr.Slider(
            minimum=16, maximum=256, value=96, step=16, label="New tokens"
        )

    prompt = gr.Textbox(
        label="Prompt",
        placeholder="Write a short story about a robot learning Rust...",
        lines=5,
    )

    with gr.Accordion("Generation settings", open=False):
        with gr.Row():
            temperature = gr.Slider(
                minimum=0, maximum=1.5, value=0.8, step=0.05,
                label="Temperature (0 = greedy)",
            )
            top_p = gr.Slider(
                minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-p"
            )
            repetition = gr.Slider(
                minimum=1.0, maximum=1.5, value=1.1, step=0.05,
                label="Repetition penalty",
            )

    run = gr.Button("Generate", variant="primary")
    output = gr.Textbox(label="Model continuation", lines=12)
    inputs = [model_choice, prompt, max_tokens, temperature, top_p, repetition]
    run.click(fn=generate, inputs=inputs, outputs=output)
    prompt.submit(fn=generate, inputs=inputs, outputs=output)

    gr.Markdown(
        "**First use:** the Space downloads model files and pins the matching "
        "Talos source. Generation quality is experimental."
    )


if __name__ == "__main__":
    demo.queue().launch()




	import json
	import re
	import shutil
	import subprocess
	import sys
	from dataclasses import fields
	from pathlib import Path

	import gradio as gr
	import spaces
	import torch
	from huggingface_hub import HfApi, hf_hub_download
	from safetensors.torch import load_file


	TITLE = "Promethean Studios Model Playground"
	MODELS = {
	"Talos 10M": "PrometheanStudio/talos-10m",
	"Styx 100M": "PrometheanStudio/styx-100m",
	}


	TALOS_COMMIT = "e230fed9c18a5467e69ed0f15778a5b06d59cb5a"
	TALOS_URL = "https://github.com/Promethean-Studios/talos.git"
	CACHE = Path("/tmp/promethean-playground")
	SOURCE = CACHE / "talos"
	CACHE.mkdir(parents=True, exist_ok=True)

	MODEL_CACHE = {}
	TOKENIZER_CACHE = {}


	FALLBACK_CONFIGS = {
	"PrometheanStudio/talos-10m": {
	"vocab_size": 1024,
	"max_seq_len": 512,
	"hidden_size": 448,
	"num_layers": 3,
	"num_attention_heads": 28,
	"num_kv_heads": 14,
	"head_dim": 16,
	"ffn_type": "dense",
	"intermediate_size": 1792,
	"attention_type": "full",
	"rope_theta": 10000.0,
	"layer_norm_eps": 1e-5,
	"tie_word_embeddings": False,
	},
	"PrometheanStudio/styx-100m": {
	"vocab_size": 1024,
	"max_seq_len": 512,
	"hidden_size": 1024,
	"num_layers": 6,
	"num_attention_heads": 64,
	"num_kv_heads": 32,
	"head_dim": 16,
	"ffn_type": "dense",
	"intermediate_size": 4096,
	"attention_type": "full",
	"rope_theta": 10000.0,
	"layer_norm_eps": 1e-5,
	"tie_word_embeddings": False,
	},
	}


	def repo_files(repo_id):
	return HfApi().list_repo_files(repo_id=repo_id, repo_type="model")


	def find_tokenizer_file(files):
	candidates = [name for name in files if Path(name).name.lower() == "tokenizer.json"]
	if not candidates:
	raise FileNotFoundError("No tokenizer.json was found in the model repository.")
	candidates.sort(
	key=lambda name: (
	0 if name == "tokenizer/tokenizer.json" else 1,
	0 if name == "tokenizer.json" else 1,
	len(name),
	name,
	)
	)
	return candidates[0]


	def ensure_talos_source():

	marker = SOURCE / ".talos_commit"
	required_files = (
	"model/gpt.py",
	"model/config.py",
	"inference/generate.py",
	"tokenizer/tokenizer.py",
	"tokenizer/vocab.py",
	)
	source_is_ready = (
	marker.is_file()
	and marker.read_text(encoding="utf-8").strip() == TALOS_COMMIT
	and all((SOURCE / filename).is_file() for filename in required_files)
	)

	if not source_is_ready:
	if SOURCE.exists():
	shutil.rmtree(SOURCE)
	SOURCE.mkdir(parents=True, exist_ok=True)
	commands = [
	["git", "init", str(SOURCE)],
	["git", "-C", str(SOURCE), "remote", "add", "origin", TALOS_URL],
	["git", "-C", str(SOURCE), "fetch", "--depth", "1", "origin", TALOS_COMMIT],
	["git", "-C", str(SOURCE), "checkout", "--detach", "FETCH_HEAD"],
	]
	for command in commands:
	try:
	result = subprocess.run(
	command, check=False, capture_output=True, text=True, timeout=240
	)
	except (OSError, subprocess.TimeoutExpired) as exc:
	raise RuntimeError(f"Could not prepare Talos source: {exc}") from exc
	if result.returncode != 0:
	detail = (result.stderr or result.stdout or "No git output").strip()
	raise RuntimeError(f"Talos git checkout failed: {detail[-1200:]}")
	if not all((SOURCE / filename).is_file() for filename in required_files):
	raise RuntimeError("The pinned Talos revision is missing expected source files.")
	marker.write_text(TALOS_COMMIT, encoding="utf-8")
	print(f"Talos source pinned to {TALOS_COMMIT}")

	source_prefix = str(SOURCE.resolve())
	sys.path[:] = [
	entry for entry in sys.path
	if str(Path(entry or ".").resolve()) != source_prefix
	]
	sys.path.insert(0, source_prefix)


	for name, module in list(sys.modules.items()):
	if not (
	name in ("model", "inference", "tokenizer")
	or name.startswith(("model.", "inference.", "tokenizer."))
	):
	continue
	module_file = getattr(module, "__file__", None)
	try:
	belongs_to_source = bool(
	module_file and Path(module_file).resolve().is_relative_to(SOURCE.resolve())
	)
	except (OSError, ValueError):
	belongs_to_source = False
	if not belongs_to_source:
	del sys.modules[name]


	def load_tokenizer(repo_id, files):

	if repo_id in TOKENIZER_CACHE:
	return TOKENIZER_CACHE[repo_id]

	ensure_talos_source()
	from tokenizer.tokenizer import ByteLevelBPETokenizer

	filename = find_tokenizer_file(files)
	downloaded = hf_hub_download(repo_id=repo_id, filename=filename, repo_type="model")

	with open(downloaded, "r", encoding="utf-8") as handle:
	payload = json.load(handle)
	if not isinstance(payload, dict) or payload.get("type") != "ByteLevelBPE":
	raise RuntimeError(
	f"{repo_id}/{filename} is not the native Talos ByteLevelBPE format. "
	"A different tokenizer loader is required."
	)
	if payload.get("version") != 1:
	raise RuntimeError(
	f"Unsupported Talos tokenizer version {payload.get('version')!r} "
	f"in {repo_id}/{filename}."
	)

	tokenizer = ByteLevelBPETokenizer.from_file(str(downloaded))

	print(
	f"Tokenizer ready: {repo_id}/{filename}; "
	f"defined tokens={tokenizer.vocab_size}, "
	f"configured vocabulary={tokenizer.config.vocab_size}, "
	f"merges={tokenizer.merge_count}"
	)
	TOKENIZER_CACHE[repo_id] = tokenizer
	return tokenizer


	def config_dict_from(value):
	if isinstance(value, dict):
	for nested_key in ("model_config", "config", "model"):
	nested = value.get(nested_key)
	if isinstance(nested, dict) and any(
	key in nested for key in ("hidden_size", "num_layers", "vocab_size")
	):
	return nested
	return value
	if hasattr(value, "__dict__"):
	return vars(value)
	return None


	def config_candidates(repo_id, files, checkpoint):
	from model.config import ModelConfig

	allowed = {item.name for item in fields(ModelConfig)}
	required = {
	"vocab_size", "max_seq_len", "hidden_size", "num_layers",
	"num_attention_heads", "num_kv_heads",
	}
	candidates = []

	if "config.json" in files:
	path = hf_hub_download(repo_id=repo_id, filename="config.json", repo_type="model")
	try:
	with open(path, "r", encoding="utf-8") as handle:
	raw = json.load(handle)
	normalized = config_dict_from(raw)
	if normalized:
	clean = {key: value for key, value in normalized.items() if key in allowed}
	if required.issubset(clean):
	candidates.append(("repository config.json", clean))
	except Exception as exc:
	print(f"Could not use config.json for {repo_id}: {exc}")

	if isinstance(checkpoint, dict):
	for key in ("config", "model_config"):
	raw = config_dict_from(checkpoint.get(key))
	if raw:
	clean = {name: value for name, value in raw.items() if name in allowed}
	if required.issubset(clean):
	candidates.append((f"checkpoint {key}", clean))

	if repo_id in FALLBACK_CONFIGS:
	candidates.append(("known model preset", FALLBACK_CONFIGS[repo_id]))

	unique, seen = [], set()
	for source, cfg in candidates:
	signature = json.dumps(cfg, sort_keys=True, default=str)
	if signature not in seen:
	unique.append((source, cfg))
	seen.add(signature)
	return unique


	def read_weights(repo_id, files):
	safetensors_files = [name for name in files if name.endswith(".safetensors")]
	if safetensors_files:
	safetensors_files.sort(
	key=lambda name: (
	0 if name == "model.safetensors" else 1,
	len(name),
	name,
	)
	)
	filename = safetensors_files[0]
	path = hf_hub_download(repo_id=repo_id, filename=filename, repo_type="model")
	print(f"Loading inference weights: {repo_id}/{filename}")
	return load_file(path, device="cpu"), None

	checkpoint_files = [name for name in files if name.endswith((".pt", ".pth", ".bin"))]
	if not checkpoint_files:
	raise FileNotFoundError(f"No model weights found in {repo_id}.")

	def priority(name):
	basename = Path(name).name.lower()
	match = re.search(r"step[-_](\d+)", basename)
	step = int(match.group(1)) if match else -1
	return (
	0 if step >= 0 else 1,
	-step,
	0 if any(word in basename for word in ("final", "best", "model")) else 1,
	len(name),
	name,
	)

	checkpoint_files.sort(key=priority)
	filename = checkpoint_files[0]
	path = hf_hub_download(repo_id=repo_id, filename=filename, repo_type="model")
	print(f"Loading PyTorch checkpoint: {repo_id}/{filename}")

	try:
	checkpoint = torch.load(path, map_location="cpu", weights_only=True)
	except Exception as safe_exc:


	if repo_id not in MODELS.values():
	raise RuntimeError(
	f"Safe checkpoint loading failed: {safe_exc}. "
	"Publish a model.safetensors file instead."
	) from safe_exc
	print(f"Safe loading failed for {filename}; trying legacy checkpoint reader")
	try:
	checkpoint = torch.load(path, map_location="cpu", weights_only=False)
	except Exception as legacy_exc:
	raise RuntimeError(
	f"Could not load {filename}: safe={safe_exc}; legacy={legacy_exc}"
	) from legacy_exc

	state = checkpoint
	if isinstance(checkpoint, dict):
	for key in ("model", "model_state_dict", "state_dict", "model_state", "weights"):
	value = checkpoint.get(key)
	if isinstance(value, dict) and value and all(torch.is_tensor(v) for v in value.values()):
	state = value
	break
	if not isinstance(state, dict) or not state or not all(torch.is_tensor(v) for v in state.values()):
	raise RuntimeError(f"{filename} did not contain a model state dictionary.")
	return state, checkpoint if isinstance(checkpoint, dict) else None


	def normalize_state_dict(state):
	for prefix in ("module.", "_orig_mod."):
	if state and all(key.startswith(prefix) for key in state):
	state = {key[len(prefix):]: value for key, value in state.items()}
	return state


	def allowed_token_ids(tokenizer, model_vocab_size):


	text_ids = set(tokenizer.vocab.id_to_bytes.keys())
	valid_ids = text_ids \| {int(tokenizer.eos_id)}
	if not valid_ids or min(valid_ids) < 0 or max(valid_ids) >= model_vocab_size:
	raise RuntimeError(
	"Tokenizer contains IDs outside the model vocabulary: "
	f"largest={max(valid_ids, default=-1)}, model_vocab={model_vocab_size}"
	)


	all_ids = text_ids \| set(tokenizer.special_id_map().values())
	if min(all_ids, default=0) < 0 or max(all_ids, default=-1) >= model_vocab_size:
	raise RuntimeError(
	"Tokenizer special-token IDs exceed the model vocabulary: "
	f"largest={max(all_ids)}, model_vocab={model_vocab_size}"
	)
	return frozenset(valid_ids)


	def load_model(repo_id):
	if repo_id in MODEL_CACHE:
	return MODEL_CACHE[repo_id]

	files = repo_files(repo_id)
	ensure_talos_source()
	tokenizer = load_tokenizer(repo_id, files)

	state, checkpoint = read_weights(repo_id, files)
	state = normalize_state_dict(state)

	from model import TalosGPT
	from model.config import ModelConfig

	attempts = []
	for source, raw_config in config_candidates(repo_id, files, checkpoint):
	try:
	config = ModelConfig(**raw_config).derive()
	model = TalosGPT(config)
	model.load_state_dict(state, strict=True)
	valid_ids = allowed_token_ids(tokenizer, config.vocab_size)
	model.eval()
	print(
	f"Loaded {repo_id} from {source}: "
	f"{model.num_parameters():,} parameters, "
	f"context={config.max_seq_len}, generatable_ids={len(valid_ids)}"
	)
	MODEL_CACHE[repo_id] = (model, tokenizer, valid_ids)
	return MODEL_CACHE[repo_id]
	except Exception as exc:
	attempts.append(
	f"{source} (hidden={raw_config.get('hidden_size')}, "
	f"layers={raw_config.get('num_layers')}): "
	f"{type(exc).__name__}: {str(exc)[:700]}"
	)

	raise RuntimeError(
	f"The weights in {repo_id} did not strictly match any Talos config.\n"
	+ "\n".join(attempts)
	)


	def get_model(label):
	if label not in MODELS:
	raise ValueError(f"Unknown model selection: {label!r}")
	return load_model(MODELS[label])


	def choose_next_token(scores, history, temperature, top_p, repetition_penalty, allowed_mask):
	scores = scores.float().clone()

	if repetition_penalty > 1.0:
	for token_id in set(history[-128:]):
	if 0 <= token_id < scores.numel():
	if scores[token_id] < 0:
	scores[token_id] *= repetition_penalty
	else:
	scores[token_id] /= repetition_penalty

	scores = scores.masked_fill(~allowed_mask, float("-inf"))
	scores = torch.where(torch.isnan(scores), -torch.inf, scores)
	if not torch.isfinite(scores).any():
	raise RuntimeError("The model produced no finite logits for decodable tokens.")

	if temperature <= 0:
	return scores.argmax(dim=-1).view(1, 1)

	probabilities = torch.softmax(scores / max(float(temperature), 1e-5), dim=-1)
	if not torch.isfinite(probabilities).all():
	raise RuntimeError("Model output probabilities contain NaN or infinity.")

	sorted_probs, sorted_ids = torch.sort(probabilities, descending=True)
	cumulative = torch.cumsum(sorted_probs, dim=-1)
	remove = cumulative > float(top_p)
	remove[1:] = remove[:-1].clone()
	remove[0] = False
	sorted_probs[remove] = 0
	total = sorted_probs.sum()
	if not torch.isfinite(total) or total <= 0:
	return scores.argmax(dim=-1).view(1, 1)
	sampled_index = torch.multinomial(sorted_probs / total, num_samples=1)
	return sorted_ids.gather(0, sampled_index).view(1, 1)


	@spaces.GPU
	def generate(label, prompt, max_new_tokens, temperature, top_p, repetition_penalty):
	if not prompt or not prompt.strip():
	raise gr.Error("Enter a prompt first.")

	try:
	model, tokenizer, valid_ids = get_model(label)
	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
	model = model.to(device).eval()


	prompt_ids = tokenizer.encode(prompt)
	if not prompt_ids:
	raise gr.Error("The tokenizer returned no tokens for the prompt.")

	max_new_tokens = int(max_new_tokens)
	max_seq_len = int(model.config.max_seq_len)
	if not 1 <= max_new_tokens < max_seq_len:
	raise gr.Error(f"New-token count must be between 1 and {max_seq_len - 1}.")
	prompt_ids = prompt_ids[-(max_seq_len - max_new_tokens):]
	if min(prompt_ids) < 0 or max(prompt_ids) >= model.config.vocab_size:
	raise RuntimeError("Encoded prompt contains token IDs outside the model vocabulary.")

	input_ids = torch.tensor([prompt_ids], dtype=torch.long, device=device)
	allowed_mask = torch.zeros(model.config.vocab_size, dtype=torch.bool, device=device)
	allowed_mask[list(valid_ids)] = True
	history = list(prompt_ids)
	generated_ids = []

	from inference.generate import prefill, decode_step

	with torch.inference_mode():
	logits, cache = prefill(model, input_ids)
	for step in range(max_new_tokens):
	scores = logits[0, -1, :]
	next_token = choose_next_token(
	scores,
	history,
	float(temperature),
	float(top_p),
	float(repetition_penalty),
	allowed_mask,
	)
	token_id = int(next_token.item())
	if token_id == tokenizer.eos_id:
	break
	generated_ids.append(token_id)
	history.append(token_id)
	if step + 1 < max_new_tokens:

	logits = decode_step(model, cache, next_token, position=cache.length)

	return tokenizer.decode(generated_ids)

	except gr.Error:
	raise
	except Exception as exc:
	print(f"Generation failure for {label}: {type(exc).__name__}: {exc}")
	raise gr.Error(f"Could not load or run {label}: {type(exc).__name__}: {exc}") from exc


	with gr.Blocks(title=TITLE) as demo:
	gr.Markdown(
	"# Promethean Studios Model Playground\n"
	"Experiment with Talos and Styx checkpoints. These are research base models, "
	"not guaranteed to behave like instruction-tuned chat assistants."
	)

	with gr.Row():
	model_choice = gr.Dropdown(
	choices=list(MODELS.keys()), value="Styx 100M", label="Model"
	)
	max_tokens = gr.Slider(
	minimum=16, maximum=256, value=96, step=16, label="New tokens"
	)

	prompt = gr.Textbox(
	label="Prompt",
	placeholder="Write a short story about a robot learning Rust...",
	lines=5,
	)

	with gr.Accordion("Generation settings", open=False):
	with gr.Row():
	temperature = gr.Slider(
	minimum=0, maximum=1.5, value=0.8, step=0.05,
	label="Temperature (0 = greedy)",
	)
	top_p = gr.Slider(
	minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-p"
	)
	repetition = gr.Slider(
	minimum=1.0, maximum=1.5, value=1.1, step=0.05,
	label="Repetition penalty",
	)

	run = gr.Button("Generate", variant="primary")
	output = gr.Textbox(label="Model continuation", lines=12)
	inputs = [model_choice, prompt, max_tokens, temperature, top_p, repetition]
	run.click(fn=generate, inputs=inputs, outputs=output)
	prompt.submit(fn=generate, inputs=inputs, outputs=output)

	gr.Markdown(
	"First use: the Space downloads model files and pins the matching "
	"Talos source. Generation quality is experimental."
	)


	if __name__ == "__main__":
	demo.queue().launch()