Upload infer_bytefast60m.py

infer_bytefast60m.py

@@ -0,0 +1,940 @@
+#!/usr/bin/env python3
+"""
+infer_bytefast60m.py
+
+Correct standalone inference utility for the custom FastDeepHybridLM defined
+by bytefalcon_fast60m.py.
+
+This does NOT instantiate Falcon-H1 or any Hugging Face AutoModel class.
+It imports the exact training architecture and calls its load_model_bundle(),
+which reconstructs Fast60MConfig + FastDeepHybridLM and strictly loads model.pt.
+
+Expected checkpoint:
+    step-XXXXXXXX/
+        config.json
+        model.pt
+        tokenizer.json
+        tokenizer_config.json
+        ...
+
+Rewrite training format:
+    instruction
+
+    "source text"
+
+    "target output"<eos>
+
+For inference, rewrite mode supplies the opening output quote and lets the
+model generate the target text, closing quote, and EOS.
+"""
+
+from __future__ import annotations
+
+import argparse
+import contextlib
+import importlib.util
+import json
+import os
+import re
+import sys
+import time
+from pathlib import Path
+from types import ModuleType
+from typing import Any
+
+
+CONTEXT_LENGTH = 4096
+
+# Match the training runtime setup before importing the architecture module.
+os.environ.setdefault("USE_HUB_KERNELS", "NO")
+os.environ.setdefault("PYTORCH_ALLOC_CONF", "expandable_segments:True")
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+os.environ.setdefault("USE_ROCM_CK_GEMM", "1")
+os.environ.pop("PYTORCH_HIP_ALLOC_CONF", None)
+
+
+def is_checkpoint(path: Path) -> bool:
+    return (
+        path.is_dir()
+        and (path / "config.json").is_file()
+        and (path / "model.pt").is_file()
+    )
+
+
+def checkpoint_rank(path: Path) -> tuple[int, float, str]:
+    matches = re.findall(r"\d+", path.name)
+    step = int(matches[-1]) if matches else -1
+    try:
+        modified = path.stat().st_mtime
+    except OSError:
+        modified = 0.0
+    return step, modified, path.name
+
+
+def resolve_checkpoint(
+    path: Path,
+    *,
+    extra_bases: list[Path] | None = None,
+) -> Path:
+    """
+    Accept an exact checkpoint, a run directory, or its checkpoints directory.
+
+    Relative paths are searched from:
+      1. the current working directory;
+      2. the inference script directory;
+      3. any supplied extra bases, such as the architecture script directory.
+
+    Preference within each candidate:
+      exact directory -> final/ -> newest immediate checkpoint ->
+      newest run/checkpoints checkpoint -> initial/
+    """
+    raw_path = path.expanduser()
+    bases = [
+        Path.cwd(),
+        Path(__file__).resolve().parent,
+    ]
+    if extra_bases:
+        bases.extend(base.expanduser().resolve() for base in extra_bases)
+
+    candidate_roots: list[Path] = []
+    if raw_path.is_absolute():
+        candidate_roots.append(raw_path.resolve())
+    else:
+        candidate_roots.extend(
+            (base / raw_path).resolve()
+            for base in bases
+        )
+
+    candidate_roots = list(dict.fromkeys(candidate_roots))
+    inspected: list[dict[str, Any]] = []
+
+    def inspect_root(root: Path) -> Path | None:
+        inspected.append(
+            {
+                "candidate_root": str(root),
+                "exists": root.exists(),
+                "is_directory": root.is_dir(),
+            }
+        )
+
+        if is_checkpoint(root):
+            return root
+
+        final = root / "final"
+        if is_checkpoint(final):
+            return final
+
+        direct = (
+            sorted(
+                (
+                    child
+                    for child in root.iterdir()
+                    if child.is_dir() and is_checkpoint(child)
+                ),
+                key=checkpoint_rank,
+            )
+            if root.is_dir()
+            else []
+        )
+        if direct:
+            return direct[-1]
+
+        checkpoint_root = root / "checkpoints"
+        nested = (
+            sorted(
+                (
+                    child
+                    for child in checkpoint_root.iterdir()
+                    if child.is_dir() and is_checkpoint(child)
+                ),
+                key=checkpoint_rank,
+            )
+            if checkpoint_root.is_dir()
+            else []
+        )
+        if nested:
+            return nested[-1]
+
+        initial = root / "initial"
+        if is_checkpoint(initial):
+            print(
+                "WARNING: using initial/; this is an untrained model.",
+                file=sys.stderr,
+            )
+            return initial
+
+        for directory in (root, checkpoint_root):
+            if not directory.is_dir():
+                continue
+            for child in sorted(directory.iterdir()):
+                if not child.is_dir():
+                    continue
+                inspected.append(
+                    {
+                        "path": str(child),
+                        "has_config": (child / "config.json").is_file(),
+                        "has_model_pt": (child / "model.pt").is_file(),
+                        "files": sorted(
+                            item.name
+                            for item in child.iterdir()
+                            if item.is_file()
+                        )[:50],
+                    }
+                )
+
+        return None
+
+    for candidate_root in candidate_roots:
+        resolved = inspect_root(candidate_root)
+        if resolved is not None:
+            print(
+                f"Resolved model path from {candidate_root}",
+                file=sys.stderr,
+            )
+            return resolved
+
+    raise FileNotFoundError(
+        "Could not find a FastDeepHybridLM checkpoint containing both "
+        "config.json and model.pt.\n"
+        "The supplied --model path was searched relative to the working "
+        "directory, inference-script directory, and architecture-script "
+        "directory.\n"
+        + json.dumps(inspected, indent=2)
+        + "\n\nCurrent working directory: "
+        + str(Path.cwd())
+        + "\nInference script directory: "
+        + str(Path(__file__).resolve().parent)
+    )
+
+
+def find_architecture_script(explicit: Path | None) -> Path:
+    if explicit is not None:
+        path = explicit.expanduser().resolve()
+        if not path.is_file():
+            raise FileNotFoundError(
+                f"Architecture script does not exist: {path}"
+            )
+        return path
+
+    here = Path(__file__).resolve().parent
+    cwd = Path.cwd()
+    candidates = [
+        cwd / "bytefalcon_fast60m.py",
+        cwd / "bytefalcon.py",
+        here / "bytefalcon_fast60m.py",
+        here / "bytefalcon.py",
+    ]
+
+    for candidate in candidates:
+        if not candidate.is_file():
+            continue
+        source = candidate.read_text(
+            encoding="utf-8",
+            errors="replace",
+        )
+        required = (
+            "class Fast60MConfig",
+            "def create_model_classes",
+            "def load_model_bundle",
+        )
+        if all(marker in source for marker in required):
+            return candidate.resolve()
+
+    raise FileNotFoundError(
+        "Could not locate the custom architecture script. Pass it explicitly:\n"
+        "  --architecture-script /path/to/bytefalcon_fast60m.py"
+    )
+
+
+def load_architecture_module(path: Path) -> ModuleType:
+    module_name = "_bytefast60m_architecture"
+    specification = importlib.util.spec_from_file_location(
+        module_name,
+        path,
+    )
+    if specification is None or specification.loader is None:
+        raise RuntimeError(
+            f"Could not create an import specification for {path}"
+        )
+
+    module = importlib.util.module_from_spec(specification)
+    # Dataclasses and some runtime machinery expect the module to be present.
+    sys.modules[module_name] = module
+    specification.loader.exec_module(module)
+
+    required = (
+        "Fast60MConfig",
+        "create_model_classes",
+        "import_training_stack",
+        "load_model_bundle",
+        "load_tokenizer",
+    )
+    missing = [
+        name for name in required if not hasattr(module, name)
+    ]
+    if missing:
+        raise RuntimeError(
+            f"{path} is not the FastDeepHybridLM training script; "
+            f"missing definitions: {missing}"
+        )
+
+    return module
+
+
+def resolve_tokenizer(
+    checkpoint: Path,
+    explicit: Path | None,
+) -> Path:
+    candidates: list[Path] = []
+
+    if explicit is not None:
+        candidates.append(explicit.expanduser().resolve())
+
+    candidates.extend(
+        [
+            checkpoint,
+            checkpoint / "tokenizer",
+        ]
+    )
+
+    project = Path(__file__).resolve().parent
+    candidates.extend(
+        [
+            project / "artifacts" / "byte-tokenizer",
+            Path.cwd() / "artifacts" / "byte-tokenizer",
+        ]
+    )
+
+    for parent in list(checkpoint.parents)[:5]:
+        candidates.extend(
+            [
+                parent / "artifacts" / "byte-tokenizer",
+                parent / "byte-tokenizer",
+            ]
+        )
+
+    candidates = list(dict.fromkeys(candidates))
+    for candidate in candidates:
+        if (
+            candidate.is_dir()
+            and (
+                (candidate / "tokenizer.json").is_file()
+                or (candidate / "tokenizer.model").is_file()
+            )
+        ):
+            return candidate
+
+    raise FileNotFoundError(
+        "Tokenizer not found. Pass --tokenizer explicitly. Checked:\n"
+        + "\n".join(f"  - {path}" for path in candidates)
+    )
+
+
+def rewrite_prompt(instruction: str, source_text: str) -> str:
+    instruction = instruction.strip()
+    source = '"' + source_text + '"'
+    if instruction:
+        return instruction + "\n\n" + source + '\n\n"'
+    return source + '\n\n"'
+
+
+def control_token_id_map(
+    tokenizer: Any,
+    architecture: ModuleType,
+) -> dict[str, int]:
+    control_tokens = getattr(
+        architecture,
+        "CONTROL_TOKENS",
+        [
+            "<pad>",
+            "<bos>",
+            "<eos>",
+            "<unk>",
+            "<instruction>",
+            "<text>",
+            "<output>",
+            "<record>",
+            "<byte_start>",
+            "<byte_end>",
+        ],
+    )
+    result = {}
+    for token in control_tokens:
+        token_id = tokenizer.convert_tokens_to_ids(token)
+        if token_id is None:
+            continue
+        token_id = int(token_id)
+        if token_id >= 0:
+            result[token] = token_id
+    return result
+
+
+def blocked_generation_ids(
+    tokenizer: Any,
+    architecture: ModuleType,
+) -> list[int]:
+    mapping = control_token_id_map(tokenizer, architecture)
+    return sorted(
+        token_id
+        for token, token_id in mapping.items()
+        if token != "<eos>"
+    )
+
+
+def apply_repetition_penalty(
+    torch: Any,
+    logits: Any,
+    input_ids: Any,
+    penalty: float,
+) -> Any:
+    if penalty == 1.0:
+        return logits
+
+    used = torch.unique(input_ids)
+    selected = logits[:, used]
+    logits[:, used] = torch.where(
+        selected < 0,
+        selected * penalty,
+        selected / penalty,
+    )
+    return logits
+
+
+def sample_next_token(
+    torch: Any,
+    logits: Any,
+    *,
+    temperature: float,
+    top_k: int,
+    top_p: float,
+) -> Any:
+    if temperature <= 0:
+        return logits.argmax(dim=-1, keepdim=True)
+
+    logits = logits / max(temperature, 1e-5)
+
+    if top_k > 0:
+        top_k = min(top_k, logits.shape[-1])
+        threshold = torch.topk(
+            logits,
+            top_k,
+            dim=-1,
+        ).values[:, -1:]
+        logits = logits.masked_fill(
+            logits < threshold,
+            -float("inf"),
+        )
+
+    probabilities = torch.softmax(logits, dim=-1)
+
+    if top_p < 1.0:
+        sorted_probabilities, sorted_indices = torch.sort(
+            probabilities,
+            descending=True,
+            dim=-1,
+        )
+        cumulative = sorted_probabilities.cumsum(dim=-1)
+        remove = cumulative > top_p
+        remove[:, 1:] = remove[:, :-1].clone()
+        remove[:, 0] = False
+        sorted_probabilities = sorted_probabilities.masked_fill(
+            remove,
+            0.0,
+        )
+        denominator = sorted_probabilities.sum(
+            dim=-1,
+            keepdim=True,
+        ).clamp_min(1e-12)
+        sorted_probabilities = (
+            sorted_probabilities / denominator
+        )
+        sampled = torch.multinomial(
+            sorted_probabilities,
+            num_samples=1,
+        )
+        return sorted_indices.gather(-1, sampled)
+
+    return torch.multinomial(probabilities, num_samples=1)
+
+
+def clean_completion(text: str, rewrite_mode: bool) -> str:
+    for marker in ("<eos>", "<record>", "<pad>"):
+        position = text.find(marker)
+        if position >= 0:
+            text = text[:position]
+
+    if rewrite_mode:
+        text = text.rstrip()
+        if text.endswith('"'):
+            text = text[:-1]
+
+    return text
+
+
+def generate(
+    *,
+    architecture: ModuleType,
+    checkpoint: Path,
+    tokenizer_path: Path,
+    prompt: str,
+    max_new_tokens: int,
+    temperature: float,
+    top_k: int,
+    top_p: float,
+    repetition_penalty: float,
+    seed: int,
+    compile_model: bool,
+    compile_mode: str,
+    stream: bool,
+    rewrite_mode: bool,
+    allow_control_tokens: bool,
+    show_top_tokens: int,
+) -> tuple[str, dict[str, Any]]:
+    (
+        _np,
+        torch,
+        nn,
+        F,
+        _DataLoader,
+        _Dataset,
+    ) = architecture.import_training_stack()
+
+    if not torch.cuda.is_available():
+        raise RuntimeError(
+            "ROCm PyTorch did not expose the AMD GPU through torch.cuda."
+        )
+
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+    device = torch.device("cuda")
+    tokenizer = architecture.load_tokenizer(tokenizer_path)
+    model = architecture.load_model_bundle(
+        checkpoint,
+        torch,
+        nn,
+        F,
+    )
+    model.to(device)
+    model.eval()
+
+    blocked_ids = (
+        []
+        if allow_control_tokens
+        else blocked_generation_ids(tokenizer, architecture)
+    )
+    blocked_tensor = (
+        torch.tensor(
+            blocked_ids,
+            device=device,
+            dtype=torch.long,
+        )
+        if blocked_ids
+        else None
+    )
+
+    active_model = model
+    if compile_model:
+        active_model = torch.compile(
+            model,
+            mode=compile_mode,
+            fullgraph=False,
+            dynamic=False,
+        )
+
+    encoded = tokenizer(
+        prompt,
+        add_special_tokens=False,
+        return_tensors="pt",
+        return_token_type_ids=False,
+    )
+    input_ids = encoded.input_ids.to(device)
+
+    prompt_tokens = int(input_ids.shape[1])
+    maximum_context = int(
+        model.config.max_position_embeddings
+    )
+    if prompt_tokens >= maximum_context:
+        raise ValueError(
+            f"Prompt has {prompt_tokens} tokens and exceeds the "
+            f"{maximum_context}-token context."
+        )
+
+    max_new_tokens = min(
+        max_new_tokens,
+        maximum_context - prompt_tokens,
+    )
+
+    eos_id = int(tokenizer.eos_token_id)
+    generated_ids: list[int] = []
+
+    torch.cuda.synchronize()
+    torch.cuda.reset_peak_memory_stats()
+    started = time.perf_counter()
+
+    # This architecture has no KV/conv recurrent inference cache. It therefore
+    # recomputes the active prefix each step, matching the original CLI.
+    with torch.inference_mode():
+        for _ in range(max_new_tokens):
+            model_input = input_ids[
+                :, -maximum_context:
+            ]
+
+            with torch.autocast(
+                device_type="cuda",
+                dtype=torch.bfloat16,
+                enabled=True,
+            ):
+                logits = active_model(
+                    input_ids=model_input,
+                    return_last_logits=True,
+                ).logits[:, -1, :]
+
+            if not bool(torch.isfinite(logits).all().item()):
+                print(
+                    "Non-finite BF16 logits; retrying this token in FP32.",
+                    file=sys.stderr,
+                )
+                with torch.autocast(
+                    device_type="cuda",
+                    enabled=False,
+                ):
+                    logits = model(
+                        input_ids=model_input,
+                        return_last_logits=True,
+                    ).logits[:, -1, :].float()
+
+            if not bool(torch.isfinite(logits).all().item()):
+                logits = torch.nan_to_num(
+                    logits,
+                    nan=-float("inf"),
+                    posinf=1e4,
+                    neginf=-1e4,
+                )
+
+            if show_top_tokens > 0:
+                top_values, top_indices = torch.topk(
+                    logits,
+                    min(show_top_tokens, logits.shape[-1]),
+                    dim=-1,
+                )
+                report = [
+                    {
+                        "id": int(token_id),
+                        "token": tokenizer.decode(
+                            [int(token_id)],
+                            skip_special_tokens=False,
+                            clean_up_tokenization_spaces=False,
+                        ),
+                        "logit": float(value),
+                    }
+                    for token_id, value in zip(
+                        top_indices[0].tolist(),
+                        top_values[0].float().tolist(),
+                    )
+                ]
+                print(
+                    "raw top tokens: "
+                    + json.dumps(report, ensure_ascii=False),
+                    file=sys.stderr,
+                )
+
+            if blocked_tensor is not None:
+                logits.index_fill_(
+                    1,
+                    blocked_tensor,
+                    -float("inf"),
+                )
+
+            logits = apply_repetition_penalty(
+                torch,
+                logits,
+                model_input,
+                repetition_penalty,
+            )
+
+            if not bool(torch.isfinite(logits).any().item()):
+                next_token = torch.tensor(
+                    [[int(tokenizer.eos_token_id)]],
+                    device=device,
+                    dtype=torch.long,
+                )
+            else:
+                next_token = sample_next_token(
+                    torch,
+                    logits,
+                    temperature=temperature,
+                    top_k=top_k,
+                    top_p=top_p,
+                )
+
+            token_id = int(next_token.item())
+            if token_id in blocked_ids:
+                raise RuntimeError(
+                    "A reserved control token escaped masking: "
+                    f"id={token_id}, token={tokenizer.decode([token_id], skip_special_tokens=False)!r}"
+                )
+            generated_ids.append(token_id)
+            input_ids = torch.cat(
+                (input_ids, next_token),
+                dim=-1,
+            )
+
+            if stream:
+                piece = tokenizer.decode(
+                    [token_id],
+                    skip_special_tokens=False,
+                    clean_up_tokenization_spaces=False,
+                )
+                print(piece, end="", flush=True)
+
+            if token_id == eos_id:
+                break
+
+    torch.cuda.synchronize()
+    elapsed = time.perf_counter() - started
+
+    raw_completion = tokenizer.decode(
+        generated_ids,
+        skip_special_tokens=False,
+        clean_up_tokenization_spaces=False,
+    )
+    completion = clean_completion(
+        raw_completion,
+        rewrite_mode,
+    )
+
+    if stream:
+        print()
+
+    metrics = {
+        "architecture": model.config.architecture,
+        "model_type": model.config.model_type,
+        "checkpoint": str(checkpoint),
+        "tokenizer": str(tokenizer_path),
+        "parameters": sum(
+            parameter.numel()
+            for parameter in model.parameters()
+        ),
+        "prompt_tokens": prompt_tokens,
+        "generated_tokens": len(generated_ids),
+        "elapsed_seconds": elapsed,
+        "tokens_per_second": (
+            len(generated_ids) / elapsed
+            if elapsed > 0
+            else None
+        ),
+        "peak_vram_gib": (
+            torch.cuda.max_memory_allocated() / (1024**3)
+        ),
+        "compiled": compile_model,
+        "blocked_control_token_ids": blocked_ids,
+        "note": (
+            "Generation recomputes the active prefix because this custom "
+            "architecture does not implement an incremental inference cache."
+        ),
+    }
+
+    return completion, metrics
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description=(
+            "Inference for the custom byte-deep-hybrid FastDeepHybridLM."
+        ),
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+
+    parser.add_argument(
+        "--model",
+        type=Path,
+        default=Path("runs/bytefast-60m"),
+        help=(
+            "Exact checkpoint, run directory, or checkpoints directory. "
+            "Relative paths are searched from the shell, script, and "
+            "architecture-script directories."
+        ),
+    )
+    parser.add_argument(
+        "--architecture-script",
+        type=Path,
+        help=(
+            "Path to bytefalcon_fast60m.py. Automatically discovered "
+            "when omitted."
+        ),
+    )
+    parser.add_argument(
+        "--tokenizer",
+        type=Path,
+        help=(
+            "Tokenizer directory. The checkpoint tokenizer is preferred."
+        ),
+    )
+
+    input_group = parser.add_mutually_exclusive_group(required=True)
+    input_group.add_argument(
+        "--prompt",
+        help="Raw language-model prompt.",
+    )
+    input_group.add_argument(
+        "--text",
+        help="Source text for rewrite mode.",
+    )
+    parser.add_argument(
+        "--instruction",
+        default="Rewrite this clearly and naturally.",
+        help="Instruction used with --text.",
+    )
+
+    parser.add_argument("--max-new-tokens", type=int, default=128)
+    parser.add_argument("--temperature", type=float, default=0.7)
+    parser.add_argument("--top-p", type=float, default=0.95)
+    parser.add_argument("--top-k", type=int, default=50)
+    parser.add_argument(
+        "--repetition-penalty",
+        type=float,
+        default=1.1,
+    )
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument(
+        "--stream",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+    )
+    parser.add_argument(
+        "--metrics",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+    )
+    parser.add_argument("--compile", action="store_true")
+    parser.add_argument(
+        "--compile-mode",
+        choices=[
+            "default",
+            "reduce-overhead",
+            "max-autotune",
+        ],
+        default="reduce-overhead",
+    )
+    parser.add_argument(
+        "--show-prompt",
+        action="store_true",
+    )
+    parser.add_argument(
+        "--allow-control-tokens",
+        action="store_true",
+        help="Allow structural tokens such as <pad>; disabled by default.",
+    )
+    parser.add_argument(
+        "--show-top-tokens",
+        type=int,
+        default=0,
+        help="Print the raw top-N logits before control-token masking.",
+    )
+
+    return parser
+
+
+def validate_args(args: argparse.Namespace) -> None:
+    if args.max_new_tokens <= 0:
+        raise ValueError("--max-new-tokens must be positive.")
+    if args.temperature < 0:
+        raise ValueError("--temperature cannot be negative.")
+    if not 0 < args.top_p <= 1:
+        raise ValueError("--top-p must be in (0, 1].")
+    if args.top_k < 0:
+        raise ValueError("--top-k cannot be negative.")
+    if args.repetition_penalty <= 0:
+        raise ValueError(
+            "--repetition-penalty must be positive."
+        )
+    if args.show_top_tokens < 0:
+        raise ValueError("--show-top-tokens must be non-negative.")
+
+
+def main() -> int:
+    args = build_parser().parse_args()
+    validate_args(args)
+
+    architecture_path = find_architecture_script(
+        args.architecture_script
+    )
+    checkpoint = resolve_checkpoint(
+        args.model,
+        extra_bases=[architecture_path.parent],
+    )
+    architecture = load_architecture_module(
+        architecture_path
+    )
+    tokenizer_path = resolve_tokenizer(
+        checkpoint,
+        args.tokenizer,
+    )
+
+    rewrite_mode = args.text is not None
+    prompt = (
+        rewrite_prompt(args.instruction, args.text)
+        if rewrite_mode
+        else args.prompt
+    )
+    assert prompt is not None
+
+    print(
+        json.dumps(
+            {
+                "checkpoint": str(checkpoint),
+                "architecture_script": str(architecture_path),
+                "tokenizer": str(tokenizer_path),
+                "rewrite_mode": rewrite_mode,
+            },
+            indent=2,
+        ),
+        file=sys.stderr,
+    )
+
+    if args.show_prompt:
+        print(
+            "----- PROMPT -----\n"
+            + prompt
+            + "\n----- END PROMPT -----",
+            file=sys.stderr,
+        )
+
+    completion, metrics = generate(
+        architecture=architecture,
+        checkpoint=checkpoint,
+        tokenizer_path=tokenizer_path,
+        prompt=prompt,
+        max_new_tokens=args.max_new_tokens,
+        temperature=args.temperature,
+        top_k=args.top_k,
+        top_p=args.top_p,
+        repetition_penalty=args.repetition_penalty,
+        seed=args.seed,
+        compile_model=args.compile,
+        compile_mode=args.compile_mode,
+        stream=args.stream,
+        rewrite_mode=rewrite_mode,
+        allow_control_tokens=args.allow_control_tokens,
+        show_top_tokens=args.show_top_tokens,
+    )
+
+    if not args.stream:
+        print(completion)
+
+    if args.metrics:
+        print(
+            "\n" + json.dumps(metrics, indent=2),
+            file=sys.stderr,
+        )
+
+    return 0
+
+
+if __name__ == "__main__":
+    try:
+        raise SystemExit(main())
+    except KeyboardInterrupt:
+        print("\nInterrupted.", file=sys.stderr)
+        raise SystemExit(130)