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