import gc import os import tempfile from collections.abc import Generator from functools import wraps from pathlib import Path from threading import Lock, Thread import spaces import torch import transformers from fastapi.responses import HTMLResponse from gradio import Server from transformers import AutoModelForCausalLM, AutoTokenizer, LogitsProcessor, LogitsProcessorList, TextIteratorStreamer print(f"Transformers version: {transformers.__version__}") HF_TOKEN = os.environ.get("HF_TOKEN") # Available model variants. The front-end populates its selector from the # /models endpoint, so adding a new entry here is enough to offer it. # Only one variant is kept loaded at a time; switching unloads the previous. MODEL_VARIANTS = { "Nalu 1.0 (Base)": {"repo_id": "khtsly/Nalu-1.0-Base", "tokenizer": "khtsly/luau-coder-1.0-preview-tokenizer", "enabled": True}, "Moana 1.0 (Base)": {"repo_id": "khtsly/Moana-1.0-Base", "tokenizer": "khtsly/luau-coder-1.0-preview-tokenizer", "enabled": True}, "Moana 1.5": {"repo_id": "khtsly/Moana-1.5-Base", "tokenizer": "khtsly/luau-coder-1.5-tokenizer", "enabled": False}, } DEFAULT_VARIANT = "Nalu 1.0 (Base)" def _repair_checkpoint_keys(weights_path: Path) -> bool: """Remove an obsolete composite-model prefix before Transformers loads.""" from safetensors import safe_open with safe_open(weights_path, framework="pt", device="cpu") as f: keys = list(f.keys()) prefixed = [key for key in keys if key.startswith("language_model.")] if not prefixed: return False if len(prefixed) != len(keys): raise RuntimeError( f"Refusing to repair mixed checkpoint keys in {weights_path}: " f"{len(prefixed)} of {len(keys)} use the language_model. prefix" ) fixed_keys = [key.removeprefix("language_model.") for key in keys] if len(set(fixed_keys)) != len(fixed_keys): raise RuntimeError(f"Prefix removal would create duplicate keys in {weights_path}") from safetensors.torch import load_file, save_file print(f"Repairing {len(keys)} checkpoint keys in {weights_path} ...") state = load_file(weights_path, device="cpu") fixed = { key.removeprefix("language_model."): tensor.contiguous() for key, tensor in state.items() } tmp_path = weights_path.with_name(f".{weights_path.name}.repair.tmp") try: save_file(fixed, tmp_path) os.replace(tmp_path, weights_path) finally: tmp_path.unlink(missing_ok=True) print("Checkpoint key repair complete") return True def _prepare_model_source(variant: str) -> str: """Return a local model directory whose safetensors keys are loadable.""" repo_id = MODEL_VARIANTS[variant]["repo_id"] cache_dir = Path(tempfile.gettempdir()) / f"model-{variant}" # Local override for the default variant (checked-in model.safetensors). if variant == DEFAULT_VARIANT: app_dir = Path(__file__).resolve().parent weights_path = app_dir / "model.safetensors" if weights_path.is_file(): _repair_checkpoint_keys(weights_path) return str(app_dir) from huggingface_hub import snapshot_download snapshot_download( repo_id=repo_id, local_dir=cache_dir, token=HF_TOKEN, allow_patterns=["*.json", "*.py", "*.jinja", "*.safetensors"], ) weights_path = cache_dir / "model.safetensors" if not weights_path.is_file(): raise FileNotFoundError(f"No model.safetensors found in {cache_dir}") _repair_checkpoint_keys(weights_path) return str(cache_dir) _model = None _tokenizer = None _current_variant = None _model_lock = Lock() def _unload_current_model() -> None: """Free the currently loaded variant so the next one can take its place.""" global _model, _current_variant, _tokenizer if _model is not None: print(f"Unloading model variant: {_current_variant}") del _model _model = None if _tokenizer is not None: print(f"Unloading tokenizer variant: {_current_variant}") del _tokenizer _tokenizer = None _current_variant = None gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() def ensure_model(variant: str): """Load `variant` on demand, unloading whatever is currently loaded.""" global _model, _current_variant, _tokenizer if variant not in MODEL_VARIANTS: print(f"Unknown model variant {variant!r}; falling back to {DEFAULT_VARIANT}") variant = DEFAULT_VARIANT if not MODEL_VARIANTS[variant].get("enabled", True): print(f"Model variant {variant!r} is not enabled yet; falling back to {DEFAULT_VARIANT}") variant = DEFAULT_VARIANT with _model_lock: if _model is not None and _current_variant == variant: return _model, _tokenizer _unload_current_model() model_targeted = MODEL_VARIANTS[variant] print(f"Loading model variant: {variant} ({model_targeted["repo_id"]}) | tokenizer: {model_targeted["tokenizer"]} ...") source = _prepare_model_source(variant) loaded = AutoModelForCausalLM.from_pretrained( source, trust_remote_code=True, dtype=torch.bfloat16, token=HF_TOKEN, device_map="cuda", ) tokenizer = AutoTokenizer.from_pretrained( model_targeted["tokenizer"], trust_remote_code=True, token=HF_TOKEN, ) _install_transformers_516_cache_compat(loaded) _model = loaded _current_variant = variant _tokenizer = tokenizer print(f"Model variant ready: {variant}") return _model, tokenizer def _install_transformers_516_cache_compat(loaded_model) -> None: """Handle empty generation positions with older Kimi remote code.""" inner_model = loaded_model.model original_forward = inner_model.forward @wraps(original_forward) def compatible_forward(*args, **kwargs): for name in ("cache_position", "position_ids"): value = kwargs.get(name) numel = getattr(value, "numel", None) if callable(numel) and numel() == 0: kwargs[name] = None return original_forward(*args, **kwargs) inner_model.forward = compatible_forward class PresencePenaltyLogitsProcessor(LogitsProcessor): """OpenAI-style presence penalty for HF generate (which lacks it natively). Subtracts `penalty` once from the logit of every unique token already present in the context, discouraging the model from reusing tokens. """ def __init__(self, penalty: float): self.penalty = float(penalty) def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: if self.penalty == 0: return scores for batch_idx, ids in enumerate(input_ids): for token_id in set(ids.tolist()): if 0 <= token_id < scores.shape[-1]: scores[batch_idx, token_id] -= self.penalty return scores demo = Server() @demo.api() @spaces.GPU(duration=60) def predict( message: str, history: list[list] | None = None, max_tokens: int = 32, temperature: float = 0.6, top_p: float = 0.95, top_k: int = 40, min_p: float = 0.05, repetition_penalty: float = 1.05, presence_penalty: float = 0.0, model_id: str = "Nalu 1.0 (Base)", ) -> Generator[str, None, None]: # History is accepted for front-end chat compatibility; generation uses # the current message only to preserve base-model continuation behavior. raw_input_text = message max_tokens = int(max_tokens) active_model, tokenizer = ensure_model(model_id) if next(active_model.parameters()).device.type != "cuda": active_model.to("cuda") inputs = tokenizer(raw_input_text, return_tensors="pt") token_count = inputs["input_ids"].shape[1] print(f"Generation request: {len(raw_input_text)} chars, {token_count} tokens") if token_count == 0: seed_token_id = tokenizer.bos_token_id if seed_token_id is None: candidate = tokenizer.convert_tokens_to_ids("[BOS]") if isinstance(candidate, int) and candidate >= 0: seed_token_id = candidate if seed_token_id is None: yield "Please enter a non-empty prompt." return inputs["input_ids"] = torch.tensor([[seed_token_id]], dtype=torch.long) inputs["attention_mask"] = torch.ones((1, 1), dtype=torch.long) token_count = 1 print(f"Empty prompt: seeded generation with token {seed_token_id}") inputs = inputs.to("cuda") streamer = TextIteratorStreamer(tokenizer, skip_prompt=True) generation_kwargs = dict( **inputs, streamer=streamer, max_new_tokens=int(max_tokens), pad_token_id=tokenizer.eos_token_id, use_cache=True, generation_mode=True, ) temperature = float(temperature) top_p = float(top_p) top_k = int(top_k) min_p = float(min_p) repetition_penalty = float(repetition_penalty) presence_penalty = float(presence_penalty) if temperature > 0.0: generation_kwargs.update(do_sample=True, temperature=temperature) if 0.0 < top_p < 1.0: generation_kwargs["top_p"] = top_p if top_k > 0: generation_kwargs["top_k"] = top_k if 0.0 < min_p <= 1.0: generation_kwargs["min_p"] = min_p else: generation_kwargs.update(do_sample=False) if repetition_penalty != 1.0: generation_kwargs["repetition_penalty"] = repetition_penalty if presence_penalty != 0.0: generation_kwargs["logits_processor"] = LogitsProcessorList( [PresencePenaltyLogitsProcessor(presence_penalty)] ) generation_error = [] def run_generation(): try: active_model.generate(**generation_kwargs) except Exception as exc: generation_error.append(exc) streamer.end() thread = Thread(target=run_generation) thread.start() output_buffer = "" for new_text in streamer: output_buffer += new_text yield output_buffer thread.join() if generation_error: exc = generation_error[0] message_out = f"Generation failed: {type(exc).__name__}: {exc}" print(message_out) yield f"{output_buffer}\n\n{message_out}" if output_buffer else message_out @demo.get("/models") async def list_models(): return { "default": DEFAULT_VARIANT, "variants": [ {"id": name, "enabled": bool(cfg.get("enabled", True))} for name, cfg in sorted(MODEL_VARIANTS.items()) ], } @demo.get("/", response_class=HTMLResponse) async def homepage(): html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html") with open(html_path, "r", encoding="utf-8") as f: return f.read() if __name__ == "__main__": demo.launch(show_error=True)