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