app.py

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
from tokenizers import Tokenizer

TITLE = "Promethean Studios Model Playground"

MODELS = {
    "Talos 10M": "PrometheanStudio/talos-10m",
    "Styx 100M": "PrometheanStudio/styx-100m",
}

# This is the Talos revision that produced the released model architecture.
# Do not change this to "main": the repository's current main branch has a
# different layout and is not guaranteed to load these older checkpoints.
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 = {}

# Known architecture presets are used only if a repository's config.json does
# not match its own weights. A candidate is accepted only after strict loading.
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 load_tokenizer(repo_id, files):
    if repo_id in TOKENIZER_CACHE:
        return TOKENIZER_CACHE[repo_id]

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

    try:
        with open(downloaded, "r", encoding="utf-8") as handle:
            payload = json.load(handle)
    except Exception as exc:
        raise RuntimeError(f"Could not read {filename}: {exc}") from exc

    if not isinstance(payload, dict):
        raise RuntimeError(f"{filename} must contain a JSON object.")

    # Do not blindly coerce the tokenizer schema. The previous version of
    # this app changed numeric versions into strings, which can itself trigger
    # a Rust deserialization error depending on the uploaded file/schema.
    # Try the original file first, then only try the alternate version type
    # when the parser specifically rejects the original serialization.
    tokenizer = None
    parser_errors = []
    candidates = [("original", downloaded, payload)]
    version = payload.get("version")

    if isinstance(version, str):
        try:
            numeric_version = float(version)
            if numeric_version.is_integer():
                numeric_version = int(numeric_version)
            alternate = dict(payload)
            alternate["version"] = numeric_version
            candidates.append(("numeric-version", None, alternate))
        except (ValueError, OverflowError):
            pass
    elif isinstance(version, (int, float)) and not isinstance(version, bool):
        alternate = dict(payload)
        alternate["version"] = f"{version:.1f}" if isinstance(version, float) else f"{version}.0"
        candidates.append(("string-version", None, alternate))

    for candidate_name, candidate_path, candidate_payload in candidates:
        try:
            if candidate_path is not None:
                tokenizer = Tokenizer.from_file(str(candidate_path))
            else:
                candidate_path = CACHE / (
                    f"{repo_id.replace('/', '_')}_tokenizer_{candidate_name}.json"
                )
                with open(candidate_path, "w", encoding="utf-8") as handle:
                    json.dump(candidate_payload, handle, ensure_ascii=False)
                tokenizer = Tokenizer.from_file(str(candidate_path))
            print(f"Tokenizer loaded using {candidate_name} serialization")
            break
        except Exception as exc:
            parser_errors.append(f"{candidate_name}: {type(exc).__name__}: {exc}")

    if tokenizer is None:
        preview = json.dumps(payload, ensure_ascii=False)[:240]
        raise RuntimeError(
            f"Tokenizer parser rejected {repo_id}/{filename}. "
            f"Top-level version={version!r}; beginning of parsed JSON={preview!r}. "
            + " | ".join(parser_errors)
        )

    TOKENIZER_CACHE[repo_id] = tokenizer
    print(
        f"Tokenizer ready: {repo_id}/{filename}; "
        f"vocabulary entries={tokenizer.get_vocab_size(with_added_tokens=True)}"
    )
    return tokenizer


def ensure_talos_source():
    """Check out and import the exact Talos revision used for these checkpoints."""
    marker = SOURCE / ".talos_commit"
    source_is_ready = (
        marker.exists()
        and marker.read_text(encoding="utf-8").strip() == TALOS_COMMIT
        and (SOURCE / "model" / "gpt.py").is_file()
        and (SOURCE / "model" / "config.py").is_file()
    )

    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:
            result = subprocess.run(
                command,
                check=False,
                capture_output=True,
                text=True,
                timeout=240,
            )
            if result.returncode != 0:
                details = (result.stderr or result.stdout or "No git output").strip()
                raise RuntimeError(
                    f"Could not prepare Talos source ({' '.join(command)}): {details[-1200:]}"
                )

        marker.write_text(TALOS_COMMIT, encoding="utf-8")
        print(f"Talos architecture source pinned to {TALOS_COMMIT}")

    # Do this even when source files were cached from a prior app process:
    # /tmp may survive a restart while Python's sys.path does not.
    source_prefix = str(SOURCE.resolve())
    sys.path[:] = [p for p in sys.path if str(Path(p or ".").resolve()) != source_prefix]
    sys.path.insert(0, source_prefix)

    # Discard only similarly named modules imported from somewhere else.
    # Never unload modules already coming from this pinned checkout.
    for name, module in list(sys.modules.items()):
        if not (name == "model" or name.startswith("model.") or name == "inference" or name.startswith("inference.")):
            continue
        module_file = getattr(module, "__file__", None)
        try:
            belongs_to_source = 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 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):
    # Only pass actual ModelConfig fields to the constructor. Hugging Face
    # metadata such as "model_type" is not part of Talos ModelConfig.
    from model.config import ModelConfig

    allowed = {item.name for item in fields(ModelConfig)}
    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: val for key, val in normalized.items() if key in allowed}
                required = {
                    "vocab_size", "max_seq_len", "hidden_size", "num_layers",
                    "num_attention_heads", "num_kv_heads",
                }
                if required.issubset(clean):
                    candidates.append(("repository config.json", clean))
        except Exception as exc:
            print(f"Ignoring unreadable 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: val for name, val in raw.items() if name in allowed}
                required = {
                    "vocab_size", "max_seq_len", "hidden_size", "num_layers",
                    "num_attention_heads", "num_kv_heads",
                }
                if required.issubset(clean):
                    candidates.append((f"checkpoint {key}", clean))

    fallback = FALLBACK_CONFIGS.get(repo_id)
    if fallback:
        candidates.append(("known model preset", fallback))

    # Deduplicate equivalent config dictionaries while preserving priority.
    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 Path(name).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 .safetensors or PyTorch checkpoint found in {repo_id}.")

    def checkpoint_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=checkpoint_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 TypeError:
        # Older PyTorch releases do not expose the weights_only parameter.
        checkpoint = torch.load(path, map_location="cpu")
    except Exception as safe_exc:
        # These two model IDs are fixed to the owner's own public repositories.
        # Some legacy training checkpoints carry ordinary Python config/RNG
        # objects that weights_only=True rejects. Keep safetensors as the first
        # choice; only use the legacy unpickler for these explicit repositories.
        if repo_id not in MODELS.values():
            raise RuntimeError(
                f"Safe loading of {filename} failed: {type(safe_exc).__name__}: {safe_exc}. "
                "Prefer publishing model.safetensors for inference."
            ) from safe_exc
        try:
            checkpoint = torch.load(path, map_location="cpu", weights_only=False)
        except Exception as legacy_exc:
            raise RuntimeError(
                f"Could not load {filename} (safe load failed with {type(safe_exc).__name__}: "
                f"{safe_exc}; legacy load failed with {type(legacy_exc).__name__}: {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 plain model state dictionary.")
    return state, checkpoint if isinstance(checkpoint, dict) else None


def normalize_state_dict(state):
    # Strip common wrapper prefixes only when they exist on every tensor key.
    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 load_model(repo_id):
    if repo_id in MODEL_CACHE:
        return MODEL_CACHE[repo_id]

    files = repo_files(repo_id)
    tokenizer = load_tokenizer(repo_id, files)
    # Prepare the class modules before reading a legacy .pt checkpoint, which
    # may contain a pickled ModelConfig object from this exact code revision.
    ensure_talos_source()
    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)
            model.eval()

            vocab = tokenizer.get_vocab()
            max_token_id = max(vocab.values(), default=-1)
            if max_token_id >= config.vocab_size:
                raise RuntimeError(
                    f"Tokenizer emits token id {max_token_id}, but model vocabulary has "
                    f"only {config.vocab_size} rows. This repository's tokenizer and weights are incompatible."
                )

            count = model.num_parameters()
            print(
                f"Loaded {repo_id} using {source}: {count:,} parameters; "
                f"vocab={config.vocab_size}, context={config.max_seq_len}"
            )
            MODEL_CACHE[repo_id] = (model, tokenizer)
            return model, tokenizer
        except Exception as exc:
            attempts.append(
                f"{source} ({raw_config.get('hidden_size')} hidden, "
                f"{raw_config.get('num_layers')} layers): {type(exc).__name__}: {str(exc)[:700]}"
            )

    raise RuntimeError(
        f"The weights in {repo_id} did not strictly match the available Talos architecture configs. "
        "This is a real checkpoint/config mismatch, not a Transformers model_type problem.\n"
        + "\n".join(attempts)
    )


def get_model(label):
    return load_model(MODELS[label])


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

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

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

    probabilities = torch.softmax(scores / max(float(temperature), 1e-5), dim=-1)
    sorted_probs, sorted_ids = torch.sort(probabilities, descending=True)
    cumulative = torch.cumsum(sorted_probs, dim=-1)

    # Keep the first token that crosses the nucleus threshold.
    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)
    sorted_probs = sorted_probs / total
    sampled_index = torch.multinomial(sorted_probs, 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 = get_model(label)
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model = model.to(device).eval()

        encoded = tokenizer.encode(prompt)
        if not encoded.ids:
            raise gr.Error("The tokenizer returned no tokens for that prompt.")

        max_new_tokens = int(max_new_tokens)
        max_seq_len = int(model.config.max_seq_len)
        if max_new_tokens < 1 or max_new_tokens >= max_seq_len:
            raise gr.Error(f"New-token count must be between 1 and {max_seq_len - 1}.")

        prompt_ids = encoded.ids
        # Reserve room for all requested output tokens, retaining the latest
        # prompt tokens if the input is longer than the model's context window.
        keep = max_seq_len - max_new_tokens
        prompt_ids = prompt_ids[-keep:]
        input_ids = torch.tensor([prompt_ids], dtype=torch.long, device=device)
        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),
                )
                token_id = int(next_token.item())
                generated_ids.append(token_id)
                history.append(token_id)

                if step + 1 < max_new_tokens:
                    logits = decode_step(
                        model,
                        cache,
                        next_token.to(device),
                        position=input_ids.shape[1] + step,
                    )

        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()