#!/usr/bin/env python3 from __future__ import annotations import argparse import json import os import sys import time from pathlib import Path from typing import Any # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description=( "Generate text with ObsidianSmall-Base / Multiscreen " "using full-sequence decoding." ) ) parser.add_argument( "--model-dir", type=Path, default=Path(__file__).resolve().parent, ) prompt_group = parser.add_mutually_exclusive_group() prompt_group.add_argument("--prompt", type=str) prompt_group.add_argument("--prompt-file", type=Path) prompt_group.add_argument("--interactive", action="store_true") parser.add_argument( "--backend", choices=("auto", "torch", "triton"), default="auto", ) parser.add_argument( "--device", choices=("auto", "cpu", "cuda"), default="auto", ) parser.add_argument( "--dtype", choices=("auto", "float32", "float16", "bfloat16"), default="auto", ) parser.add_argument( "--max-new-tokens", type=int, default=128, ) # More conservative defaults than before. parser.add_argument( "--temperature", type=float, default=0.7, ) parser.add_argument( "--top-k", type=int, default=32, ) parser.add_argument( "--top-p", type=float, default=0.90, ) parser.add_argument( "--min-p", type=float, default=0.0, help=( "Relative minimum probability threshold. " "0 disables min-p." ), ) parser.add_argument( "--repetition-penalty", type=float, default=1.08, help="1.0 disables repetition penalty.", ) parser.add_argument( "--presence-penalty", type=float, default=0.0, ) parser.add_argument( "--frequency-penalty", type=float, default=0.0, ) parser.add_argument( "--repeat-window", type=int, default=256, help=( "Only tokens inside this recent window affect " "repetition/frequency penalties. 0 means all." ), ) parser.add_argument( "--seed", type=int, default=1337, ) parser.add_argument( "--greedy", action="store_true", help="Equivalent to temperature=0.", ) parser.add_argument( "--completion-only", action="store_true", help="Print only newly generated text.", ) parser.add_argument( "--stream", action="store_true", help="Display newly generated text as generation proceeds.", ) parser.add_argument( "--verbose", action="store_true", ) return parser.parse_args() # --------------------------------------------------------------------------- # Checkpoint loading # --------------------------------------------------------------------------- def load_safetensors_checkpoint( model_dir: Path, ) -> tuple[dict[str, Any], str]: try: from safetensors.torch import load_file except ImportError as exc: raise RuntimeError( "safetensors is required.\n" "Install with:\n" " pip install -U safetensors" ) from exc index_path = model_dir / "model.safetensors.index.json" if index_path.is_file(): index = json.loads( index_path.read_text(encoding="utf-8") ) weight_map = index.get("weight_map") if not isinstance(weight_map, dict) or not weight_map: raise RuntimeError( f"Invalid Safetensors index: {index_path}" ) shard_names = sorted(set(weight_map.values())) state_dict: dict[str, Any] = {} print( f"Loading {len(shard_names)} " "Safetensors shards..." ) for shard_name in shard_names: shard_path = model_dir / shard_name if not shard_path.is_file(): raise FileNotFoundError( f"Missing checkpoint shard: {shard_path}" ) print(f" {shard_path.name}") shard = load_file( str(shard_path), device="cpu", ) overlap = state_dict.keys() & shard.keys() if overlap: raise RuntimeError( "Duplicate tensor names between shards: " f"{sorted(overlap)[:10]}" ) state_dict.update(shard) return state_dict, str(index_path) model_path = model_dir / "model.safetensors" if model_path.is_file(): return ( load_file( str(model_path), device="cpu", ), str(model_path), ) candidates = sorted( model_dir.glob("*.safetensors") ) candidates = [ path for path in candidates if not ( path.name.startswith("model-") and "-of-" in path.name ) ] if len(candidates) == 1: path = candidates[0] return ( load_file( str(path), device="cpu", ), str(path), ) raise FileNotFoundError( "No unambiguous Safetensors checkpoint found." ) def load_checkpoint( model_dir: Path, torch: Any, ) -> tuple[Any, str]: if ( (model_dir / "model.safetensors").is_file() or (model_dir / "model.safetensors.index.json").is_file() or any(model_dir.glob("*.safetensors")) ): return load_safetensors_checkpoint( model_dir ) pth_path = model_dir / "lit_model.pth" if pth_path.is_file(): try: payload = torch.load( pth_path, map_location="cpu", weights_only=True, ) except TypeError: payload = torch.load( pth_path, map_location="cpu", ) return payload, str(pth_path) files = "\n".join( f" {path.name}" for path in sorted(model_dir.iterdir()) if path.is_file() ) raise FileNotFoundError( "No model checkpoint found.\n\n" "Expected:\n" " model.safetensors\n" " model.safetensors.index.json\n" " lit_model.pth\n\n" f"Files in {model_dir}:\n{files}" ) # --------------------------------------------------------------------------- # State-dict compatibility # --------------------------------------------------------------------------- def unwrap_state_dict( payload: Any, torch: Any, ) -> dict[str, Any]: if isinstance(payload, dict): for container_name in ( "model", "state_dict", "model_state_dict", ): candidate = payload.get(container_name) if isinstance(candidate, dict): payload = candidate break if not isinstance(payload, dict): raise TypeError( "Unsupported checkpoint object: " f"{type(payload).__name__}" ) result = { str(name): value for name, value in payload.items() if torch.is_tensor(value) } if not result: raise RuntimeError( "Checkpoint contains no tensors." ) return result def strip_prefix( state_dict: dict[str, Any], prefix: str, ) -> dict[str, Any]: output = {} for original_name, tensor in state_dict.items(): name = original_name while name.startswith(prefix): name = name[len(prefix):] output[name] = tensor return output def normalize_state_dict( payload: Any, model: Any, torch: Any, verbose: bool, ) -> dict[str, Any]: raw = unwrap_state_dict( payload, torch, ) model_state = model.state_dict() expected = set(model_state) candidates: list[ tuple[str, dict[str, Any]] ] = [ ("raw", raw), ( "strip _orig_mod.", strip_prefix(raw, "_orig_mod."), ), ( "strip module.", strip_prefix(raw, "module."), ), ( "strip model.", strip_prefix(raw, "model."), ), ] wrappers_removed = raw for prefix in ( "_orig_mod.", "module.", ): wrappers_removed = strip_prefix( wrappers_removed, prefix, ) candidates.append( ( "strip compile/DDP wrappers", wrappers_removed, ) ) candidates.append( ( "strip wrappers + model.", strip_prefix( wrappers_removed, "model.", ), ) ) best_name = "" best_state = None best_score = None if verbose: print() print("State-dict compatibility:") for name, candidate in candidates: keys = set(candidate) matched = len( keys & expected ) missing = len( expected - keys ) unexpected = len( keys - expected ) score = ( matched * 1_000_000 - missing * 1000 - unexpected ) if verbose: print( f" {name:30s} " f"matched={matched:4d} " f"missing={missing:4d} " f"unexpected={unexpected:4d}" ) if ( best_score is None or score > best_score ): best_score = score best_name = name best_state = candidate assert best_state is not None checkpoint_keys = set(best_state) missing = sorted( expected - checkpoint_keys ) unexpected = sorted( checkpoint_keys - expected ) if missing or unexpected: raise RuntimeError( "Checkpoint does not exactly match the model.\n" f"Mapping: {best_name}\n" f"Missing: {missing}\n" f"Unexpected: {unexpected}" ) # Also check tensor SHAPES before load_state_dict(). shape_errors = [] for name, expected_tensor in model_state.items(): actual_tensor = best_state[name] if actual_tensor.shape != expected_tensor.shape: shape_errors.append( ( name, tuple(actual_tensor.shape), tuple(expected_tensor.shape), ) ) if shape_errors: lines = [ "Checkpoint tensor shape mismatch:" ] for name, got, expected_shape in shape_errors[:20]: lines.append( f" {name}: " f"checkpoint={got}, " f"model={expected_shape}" ) raise RuntimeError( "\n".join(lines) ) if verbose: print( f"Selected mapping: {best_name}" ) return best_state # --------------------------------------------------------------------------- # Sampling # --------------------------------------------------------------------------- def apply_repetition_penalties( logits: Any, previous_tokens: Any, torch: Any, repetition_penalty: float, presence_penalty: float, frequency_penalty: float, ) -> Any: if previous_tokens.numel() == 0: return logits if ( repetition_penalty == 1.0 and presence_penalty == 0.0 and frequency_penalty == 0.0 ): return logits unique_tokens, counts = torch.unique( previous_tokens, return_counts=True, ) token_logits = logits[ :, unique_tokens, ] if repetition_penalty != 1.0: token_logits = torch.where( token_logits > 0, token_logits / repetition_penalty, token_logits * repetition_penalty, ) if presence_penalty != 0.0: token_logits = ( token_logits - presence_penalty ) if frequency_penalty != 0.0: token_logits = ( token_logits - counts.to( dtype=token_logits.dtype ).unsqueeze(0) * frequency_penalty ) logits = logits.clone() logits[ :, unique_tokens, ] = token_logits return logits def sample_token( logits: Any, previous_tokens: Any, torch: Any, temperature: float, top_k: int, top_p: float, min_p: float, repetition_penalty: float, presence_penalty: float, frequency_penalty: float, generator: Any, ) -> Any: logits = apply_repetition_penalties( logits=logits, previous_tokens=previous_tokens, torch=torch, repetition_penalty=repetition_penalty, presence_penalty=presence_penalty, frequency_penalty=frequency_penalty, ) if temperature <= 0.0: return logits.argmax( dim=-1, keepdim=True, ) logits = logits / temperature # Top-k. if top_k > 0: k = min( top_k, logits.size(-1), ) cutoff = torch.topk( logits, k=k, dim=-1, ).values[:, -1:] logits = logits.masked_fill( logits < cutoff, float("-inf"), ) # Min-p. # # Keep tokens whose probability is at least: # # min_p * P(best token) # # Using logits avoids an unnecessary full softmax here. if 0.0 < min_p < 1.0: max_logits = logits.max( dim=-1, keepdim=True, ).values min_logit = ( max_logits + torch.log( torch.tensor( min_p, dtype=logits.dtype, device=logits.device, ) ) ) logits = logits.masked_fill( logits < min_logit, float("-inf"), ) # Top-p / nucleus. if 0.0 < top_p < 1.0: sorted_logits, sorted_indices = torch.sort( logits, descending=True, dim=-1, ) sorted_probs = torch.softmax( sorted_logits, dim=-1, ) cumulative = sorted_probs.cumsum( dim=-1, ) remove = cumulative > top_p # Always retain the token which crosses the threshold. remove[:, 1:] = remove[:, :-1].clone() remove[:, 0] = False sorted_logits = sorted_logits.masked_fill( remove, float("-inf"), ) filtered_logits = torch.full_like( logits, float("-inf"), ) logits = filtered_logits.scatter( dim=-1, index=sorted_indices, src=sorted_logits, ) probabilities = torch.softmax( logits, dim=-1, ) if not torch.isfinite(probabilities).all(): raise RuntimeError( "Sampling probabilities became non-finite." ) if torch.any( probabilities.sum( dim=-1 ) <= 0 ): raise RuntimeError( "Sampling removed every possible token." ) return torch.multinomial( probabilities, num_samples=1, generator=generator, ) # --------------------------------------------------------------------------- # Output extraction # --------------------------------------------------------------------------- def extract_logits( output: Any, torch: Any, ) -> Any: if torch.is_tensor(output): return output if hasattr(output, "logits"): logits = output.logits if torch.is_tensor(logits): return logits if isinstance( output, (tuple, list), ): for item in output: if torch.is_tensor(item): return item raise TypeError( "Unable to extract logits from model output: " f"{type(output).__name__}" ) # --------------------------------------------------------------------------- # Prompt # --------------------------------------------------------------------------- def resolve_prompt( args: argparse.Namespace, ) -> str: if args.prompt is not None: return args.prompt if args.prompt_file is not None: return args.prompt_file.read_text( encoding="utf-8" ) if args.interactive: return "" raise SystemExit( "Provide --prompt, --prompt-file, " "or --interactive." ) # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def main() -> None: args = parse_args() model_dir = args.model_dir.resolve() runtime_dir = model_dir / "runtime" if not model_dir.is_dir(): raise FileNotFoundError( f"Model directory does not exist: {model_dir}" ) if not runtime_dir.is_dir(): raise FileNotFoundError( f"Runtime directory does not exist: {runtime_dir}" ) # Backend must be set before importing the custom runtime. os.environ[ "MULTISCREEN_BACKEND" ] = args.backend os.environ.setdefault( "TOKENIZERS_PARALLELISM", "false", ) sys.path.insert( 0, str(runtime_dir), ) import torch from litgpt import Config, GPT, Tokenizer # ------------------------------------------------------------------ # Device # ------------------------------------------------------------------ if args.device == "auto": device = torch.device( "cuda" if torch.cuda.is_available() else "cpu" ) else: device = torch.device( args.device ) if ( device.type == "cuda" and not torch.cuda.is_available() ): raise RuntimeError( "CUDA was requested but is unavailable." ) if ( args.backend == "triton" and device.type != "cuda" ): raise RuntimeError( "The Triton backend requires CUDA." ) # ------------------------------------------------------------------ # Dtype # ------------------------------------------------------------------ if args.dtype == "auto": if device.type == "cuda": capability = torch.cuda.get_device_capability( device ) # Ampere/Ada/Hopper have native BF16 acceleration. if ( capability[0] >= 8 and torch.cuda.is_bf16_supported() ): dtype = torch.bfloat16 else: dtype = torch.float16 else: dtype = torch.float32 else: dtype = { "float32": torch.float32, "float16": torch.float16, "bfloat16": torch.bfloat16, }[args.dtype] # ------------------------------------------------------------------ # RNG # ------------------------------------------------------------------ torch.manual_seed( args.seed ) if device.type == "cuda": torch.cuda.manual_seed_all( args.seed ) generator_device = ( "cuda" if device.type == "cuda" else "cpu" ) generator = torch.Generator( device=generator_device ) generator.manual_seed( args.seed ) # ------------------------------------------------------------------ # Config # ------------------------------------------------------------------ config_path = ( model_dir / "model_config.yaml" ) if not config_path.is_file(): raise FileNotFoundError( f"Missing config: {config_path}" ) print( f"Config: {config_path}" ) config = Config.from_file( config_path ) # ------------------------------------------------------------------ # Model # ------------------------------------------------------------------ print( "Model: instantiating..." ) model = GPT(config) parameter_count = sum( p.numel() for p in model.parameters() ) payload, checkpoint_source = load_checkpoint( model_dir, torch, ) print( f"Checkpoint: {checkpoint_source}" ) state_dict = normalize_state_dict( payload, model, torch, verbose=args.verbose, ) model.load_state_dict( state_dict, strict=True, ) del payload del state_dict model = model.to( device=device, dtype=dtype, ) model.eval() # ------------------------------------------------------------------ # Tokenizer # ------------------------------------------------------------------ tokenizer = Tokenizer( model_dir ) # Basic tokenizer/model compatibility check. tokenizer_vocab_size = getattr( tokenizer, "vocab_size", None, ) if ( tokenizer_vocab_size is not None and tokenizer_vocab_size > config.padded_vocab_size ): raise RuntimeError( "Tokenizer vocabulary is larger than " "the model output vocabulary." ) # ------------------------------------------------------------------ # Runtime info # ------------------------------------------------------------------ print() print("=" * 72) print("OBSIDIAN / MULTISCREEN GENERATION") print("=" * 72) print( f"Parameters: {parameter_count:,}" ) print( f"Device: {device}" ) print( f"Backend: {args.backend}" ) print( f"Precision: {dtype}" ) print( f"Context: {config.block_size:,}" ) print( f"Vocabulary: {config.vocab_size:,}" ) if device.type == "cuda": print( "GPU: " f"{torch.cuda.get_device_name(device)}" ) temperature = ( 0.0 if args.greedy else args.temperature ) print( f"Temperature: {temperature:g}" ) print( f"Top-k / top-p: " f"{args.top_k} / {args.top_p:g}" ) print( f"Min-p: {args.min_p:g}" ) print( "Repetition: " f"{args.repetition_penalty:g}" ) print( "Decoding: full sequence" ) print("=" * 72) print() # ------------------------------------------------------------------ # Generation # ------------------------------------------------------------------ def complete( prompt: str, ) -> str: encoded = tokenizer.encode( prompt, device=device, ).reshape( 1, -1, ).long() original_prompt_tokens = ( encoded.size(1) ) if original_prompt_tokens == 0: raise ValueError( "Prompt encoded to zero tokens." ) # We generate at most max_new_tokens, but the model itself sees # only the most recent block_size tokens each forward pass. total_capacity = ( original_prompt_tokens + args.max_new_tokens ) # Preallocate rather than torch.cat() every iteration. token_buffer = torch.empty( ( 1, total_capacity, ), dtype=torch.long, device=device, ) token_buffer[ :, :original_prompt_tokens, ] = encoded current_length = ( original_prompt_tokens ) generated_count = 0 if ( original_prompt_tokens > config.block_size ): print( f"[warning] Prompt has " f"{original_prompt_tokens:,} tokens; " f"only the latest " f"{config.block_size:,} are visible." ) if device.type == "cuda": torch.cuda.synchronize() start_time = time.perf_counter() last_stream_text = "" with torch.inference_mode(): for _ in range( args.max_new_tokens ): context_start = max( 0, current_length - config.block_size, ) context = token_buffer[ :, context_start:current_length, ] output = model( context ) logits = extract_logits( output, torch, ) next_logits = logits[ :, -1, :, ].float() if args.repeat_window > 0: repeat_start = max( 0, current_length - args.repeat_window, ) else: repeat_start = 0 previous_tokens = token_buffer[ 0, repeat_start:current_length, ] next_token = sample_token( logits=next_logits, previous_tokens=previous_tokens, torch=torch, temperature=temperature, top_k=args.top_k, top_p=args.top_p, min_p=args.min_p, repetition_penalty=args.repetition_penalty, presence_penalty=args.presence_penalty, frequency_penalty=args.frequency_penalty, generator=generator, ) token_buffer[ 0, current_length, ] = next_token.item() current_length += 1 generated_count += 1 if args.stream: current_generated = tokenizer.decode( token_buffer[ 0, original_prompt_tokens:current_length, ].detach().cpu() ) if current_generated.startswith( last_stream_text ): delta = current_generated[ len(last_stream_text): ] if delta: print( delta, end="", flush=True, ) last_stream_text = ( current_generated ) if ( tokenizer.eos_id is not None and next_token.item() == tokenizer.eos_id ): break if device.type == "cuda": torch.cuda.synchronize() elapsed = ( time.perf_counter() - start_time ) all_tokens = token_buffer[ 0, :current_length, ] generated_tokens = token_buffer[ 0, original_prompt_tokens:current_length, ] if args.stream: print() if args.completion_only: text = tokenizer.decode( generated_tokens .detach() .cpu() ) else: text = tokenizer.decode( all_tokens .detach() .cpu() ) if args.verbose: rate = ( generated_count / elapsed if elapsed > 0 else 0.0 ) print() print( f"[generation] " f"{generated_count} tokens " f"in {elapsed:.2f}s " f"({rate:.2f} tok/s)" ) if device.type == "cuda": memory = ( torch.cuda.max_memory_allocated( device ) / 2**20 ) print( f"[generation] " f"peak allocated VRAM: " f"{memory:.1f} MiB" ) return text # ------------------------------------------------------------------ # Interactive # ------------------------------------------------------------------ if args.interactive: print( "Interactive mode." ) print( "Commands: /quit, /exit" ) while True: try: prompt = input( "\nPrompt> " ) except EOFError: print() break if ( prompt.strip().lower() in { "/quit", "/exit", } ): break if not prompt.strip(): continue print() result = complete( prompt ) if not args.stream: print(result) else: prompt = resolve_prompt( args ) result = complete( prompt ) if not args.stream: print(result) if __name__ == "__main__": main()