Fixed app.py

app.py

@@ -1,3 +1,6 @@
+
+
+
 import json
 import re
 import shutil
@@ -11,18 +14,15 @@ 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"
 
+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")
@@ -32,8 +32,7 @@ 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,
@@ -73,10 +72,7 @@ def repo_files(repo_id):
 
 
 def find_tokenizer_file(files):
-    candidates = [
-        name for name in files
-        if Path(name).name.lower() == "tokenizer.json"
-    ]
+    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(
@@ -90,98 +86,26 @@ def find_tokenizer_file(files):
     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"
+    required_files = (
+        "model/gpt.py",
+        "model/config.py",
+        "inference/generate.py",
+        "tokenizer/tokenizer.py",
+        "tokenizer/vocab.py",
+    )
     source_is_ready = (
-        marker.exists()
+        marker.is_file()
         and marker.read_text(encoding="utf-8").strip() == TALOS_COMMIT
-        and (SOURCE / "model" / "gpt.py").is_file()
-        and (SOURCE / "model" / "config.py").is_file()
+        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],
@@ -189,42 +113,81 @@ def ensure_talos_source():
             ["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:]}"
+            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 architecture source pinned to {TALOS_COMMIT}")
+        print(f"Talos 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[:] = [
+        entry for entry in sys.path
+        if str(Path(entry 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.")):
+        if not (
+            name in ("model", "inference", "tokenizer")
+            or name.startswith(("model.", "inference.", "tokenizer."))
+        ):
             continue
         module_file = getattr(module, "__file__", None)
         try:
-            belongs_to_source = module_file and Path(module_file).resolve().is_relative_to(SOURCE.resolve())
+            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"):
@@ -240,11 +203,13 @@ def config_dict_from(value):
 
 
 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)}
+    required = {
+        "vocab_size", "max_seq_len", "hidden_size", "num_layers",
+        "num_attention_heads", "num_kv_heads",
+    }
     candidates = []
 
     if "config.json" in files:
@@ -254,35 +219,24 @@ def config_candidates(repo_id, files, checkpoint):
                 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",
-                }
+                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"Ignoring unreadable config.json for {repo_id}: {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: 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",
-                }
+                clean = {name: value for name, value in raw.items() if name in allowed}
                 if required.issubset(clean):
                     candidates.append((f"checkpoint {key}", clean))
 
-    fallback = FALLBACK_CONFIGS.get(repo_id)
-    if fallback:
-        candidates.append(("known model preset", fallback))
+    if repo_id in FALLBACK_CONFIGS:
+        candidates.append(("known model preset", FALLBACK_CONFIGS[repo_id]))
 
-    # Deduplicate equivalent config dictionaries while preserving priority.
-    unique = []
-    seen = set()
+    unique, seen = [], set()
     for source, cfg in candidates:
         signature = json.dumps(cfg, sort_keys=True, default=str)
         if signature not in seen:
@@ -296,7 +250,7 @@ def read_weights(repo_id, files):
     if safetensors_files:
         safetensors_files.sort(
             key=lambda name: (
-                0 if Path(name).name == "model.safetensors" else 1,
+                0 if name == "model.safetensors" else 1,
                 len(name),
                 name,
             )
@@ -306,14 +260,11 @@ def read_weights(repo_id, files):
         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"))
-    ]
+    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}.")
+        raise FileNotFoundError(f"No model weights found in {repo_id}.")
 
-    def checkpoint_priority(name):
+    def priority(name):
         basename = Path(name).name.lower()
         match = re.search(r"step[-_](\d+)", basename)
         step = int(match.group(1)) if match else -1
@@ -325,31 +276,27 @@ def read_weights(repo_id, files):
             name,
         )
 
-    checkpoint_files.sort(key=checkpoint_priority)
+    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 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."
+                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 load failed with {type(safe_exc).__name__}: "
-                f"{safe_exc}; legacy load failed with {type(legacy_exc).__name__}: {legacy_exc})."
+                f"Could not load {filename}: safe={safe_exc}; legacy={legacy_exc}"
             ) from legacy_exc
 
     state = checkpoint
@@ -360,27 +307,46 @@ def read_weights(repo_id, files):
                 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.")
+        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):
-    # 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 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)
-    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()
+    tokenizer = load_tokenizer(repo_id, files)
+
     state, checkpoint = read_weights(repo_id, files)
     state = normalize_state_dict(state)
 
@@ -393,59 +359,59 @@ def load_model(repo_id):
             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()
-
-            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}"
+                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)
-            return model, tokenizer
+            MODEL_CACHE[repo_id] = (model, tokenizer, valid_ids)
+            return MODEL_CACHE[repo_id]
         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]}"
+                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 the available Talos architecture configs. "
-        "This is a real checkpoint/config mismatch, not a Transformers model_type problem.\n"
+        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):
+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.shape[-1]:
+            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)
-
-    # Keep the first token that crosses the nucleus threshold.
     remove = cumulative > float(top_p)
     remove[1:] = remove[:-1].clone()
     remove[0] = False
@@ -453,8 +419,7 @@ def choose_next_token(scores, history, temperature, top_p, repetition_penalty):
     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)
+    sampled_index = torch.multinomial(sorted_probs / total, num_samples=1)
     return sorted_ids.gather(0, sampled_index).view(1, 1)
 
 
@@ -464,25 +429,26 @@ def generate(label, prompt, max_new_tokens, temperature, top_p, repetition_penal
         raise gr.Error("Enter a prompt first.")
 
     try:
-        model, tokenizer = get_model(label)
+        model, tokenizer, valid_ids = 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.")
+
+        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 max_new_tokens < 1 or max_new_tokens >= 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.")
 
-        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)
+        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 = []
 
@@ -498,18 +464,16 @@ def generate(label, prompt, max_new_tokens, temperature, top_p, repetition_penal
                     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.to(device),
-                        position=input_ids.shape[1] + step,
-                    )
+
+                    logits = decode_step(model, cache, next_token, position=cache.length)
 
         return tokenizer.decode(generated_ids)
 
@@ -517,9 +481,7 @@ def generate(label, prompt, max_new_tokens, temperature, top_p, repetition_penal
         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
+        raise gr.Error(f"Could not load or run {label}: {type(exc).__name__}: {exc}") from exc
 
 
 with gr.Blocks(title=TITLE) as demo:
@@ -531,16 +493,10 @@ with gr.Blocks(title=TITLE) as demo:
 
     with gr.Row():
         model_choice = gr.Dropdown(
-            choices=list(MODELS.keys()),
-            value="Styx 100M",
-            label="Model",
+            choices=list(MODELS.keys()), value="Styx 100M", label="Model"
         )
         max_tokens = gr.Slider(
-            minimum=16,
-            maximum=256,
-            value=96,
-            step=16,
-            label="New tokens",
+            minimum=16, maximum=256, value=96, step=16, label="New tokens"
         )
 
     prompt = gr.Textbox(
@@ -552,24 +508,14 @@ with gr.Blocks(title=TITLE) as demo:
     with gr.Accordion("Generation settings", open=False):
         with gr.Row():
             temperature = gr.Slider(
-                minimum=0,
-                maximum=1.5,
-                value=0.8,
-                step=0.05,
+                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",
+                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,
+                minimum=1.0, maximum=1.5, value=1.1, step=0.05,
                 label="Repetition penalty",
             )
 
@@ -580,10 +526,10 @@ with gr.Blocks(title=TITLE) as demo:
     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."
+        "**First use:** the Space downloads model files and pins the matching "
+        "Talos source. Generation quality is experimental."
     )
 
 
 if __name__ == "__main__":
-    demo.queue().launch()
+    demo.queue().launch()