src/mini_transformer/cli.py

"""Command-line entry points for mini_transformer."""

from __future__ import annotations

import argparse
import os
import subprocess
import sys
from collections.abc import Sequence
from pathlib import Path

import uvicorn
from omegaconf import DictConfig

from .apps import server as server_module
from .inference import run_inference
from .model_loader import (
    CONFIG_DIR_ENV,
    CONFIG_NAME_ENV,
    MODEL_NAME_ENV,
    MODELS_ENV,
    compose_config_from_dir,
    compose_model_config,
    ensure_models_root,
    list_model_names,
    snapshot_to_local,
)

DEFAULT_CONFIG_NAME = "config_inference"

__all__ = ["infer_main", "serve_main", "ui_main", "fetch_main"]


def _select_model(model: str | None) -> str:
    ensure_models_root()
    names = list_model_names()
    if not names:
        raise SystemExit(
            "No models found in `trained_models/`. Place a model folder with a configs/config_inference.yaml file there first."
        )
    if model and model in names:
        return model
    if model:
        raise SystemExit(f"Model '{model}' not found. Available models: {', '.join(names)}.")
    return names[0]


def infer_main(argv: Sequence[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Run mini_transformer inference")
    parser.add_argument("-t", "--input-text", help="Prompt to feed the model")
    parser.add_argument("-m", "--model", help="Model folder under trained_models/")
    parser.add_argument("--model-root", help="Override path to trained_models root")
    parser.add_argument(
        "-c", "--config-dir", help="Use Hydra configs from this directory instead of a model folder"
    )
    parser.add_argument(
        "--config-name",
        default=os.environ.get(CONFIG_NAME_ENV, DEFAULT_CONFIG_NAME),
        help="Hydra config name to compose (default: config_inference)",
    )
    parser.add_argument(
        "-o",
        "--override",
        action="append",
        default=[],
        help="Hydra override in dotlist form (e.g. generation.temperature=0.7)",
    )
    args = parser.parse_args(argv)

    overrides = args.override or []
    if args.model_root:
        os.environ[MODELS_ENV] = args.model_root

    if args.config_dir or os.environ.get(CONFIG_DIR_ENV):
        config_dir = args.config_dir or os.environ.get(CONFIG_DIR_ENV)
        if config_dir is None:
            raise SystemExit("Config directory not provided")
        cfg: DictConfig = compose_config_from_dir(
            config_dir,
            config_name=args.config_name,
            overrides=overrides,
        )
    else:
        model_name = _select_model(args.model)
        cfg = compose_model_config(
            model_name,
            config_name=args.config_name,
            overrides=overrides,
        )
    if args.input_text:
        cfg.input_text = args.input_text

    outputs = run_inference(cfg)
    print("\n".join(outputs))
    return 0


def serve_main(argv: Sequence[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Launch the FastAPI inference server")
    parser.add_argument("--host", default="0.0.0.0", help="Bind address (default: 0.0.0.0)")
    parser.add_argument("--port", type=int, default=8000, help="Port for uvicorn (default: 8000)")
    parser.add_argument("--reload", action="store_true", help="Enable auto-reload")
    parser.add_argument("-m", "--model", help="Model folder under trained_models/")
    parser.add_argument("--model-root", help="Override path to trained_models root")
    parser.add_argument(
        "-c", "--config-dir", help="Use configs from this directory instead of a model folder"
    )
    parser.add_argument(
        "--config-name",
        default=os.environ.get(CONFIG_NAME_ENV, DEFAULT_CONFIG_NAME),
        help="Hydra config name to compose",
    )
    args = parser.parse_args(argv)

    if args.model_root:
        os.environ[MODELS_ENV] = args.model_root

    if args.config_dir:
        os.environ[CONFIG_DIR_ENV] = args.config_dir
        os.environ.pop(MODEL_NAME_ENV, None)
    else:
        model_name = _select_model(args.model)
        os.environ[MODEL_NAME_ENV] = model_name
        os.environ.pop(CONFIG_DIR_ENV, None)

    os.environ[CONFIG_NAME_ENV] = args.config_name

    uvicorn.run(
        "mini_transformer.apps.server:app",
        host=args.host,
        port=args.port,
        reload=args.reload,
    )
    return 0


def ui_main(argv: Sequence[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Launch the Chainlit chat UI")
    parser.add_argument("--model-root", help="Override path to trained_models root")
    parser.add_argument("-m", "--model", help="Pre-select a model when the UI starts")
    parser.add_argument(
        "-c", "--config-dir", help="Use configs from this directory instead of trained_models"
    )
    parser.add_argument(
        "--host",
        default=os.environ.get("MINI_TRANSFORMER_UI_HOST", "127.0.0.1"),
        help="Host to bind Chainlit to (default: 127.0.0.1)",
    )
    parser.add_argument(
        "--port",
        type=int,
        default=int(os.environ.get("MINI_TRANSFORMER_UI_PORT", 8000)),
        help="Port for Chainlit (default: 8000)",
    )
    parser.add_argument(
        "--config-name",
        default=os.environ.get(CONFIG_NAME_ENV, DEFAULT_CONFIG_NAME),
        help="Hydra config name to compose",
    )
    args = parser.parse_args(argv)

    env = os.environ.copy()
    if args.model_root:
        env[MODELS_ENV] = args.model_root
    if args.config_dir:
        env[CONFIG_DIR_ENV] = args.config_dir
    if args.model:
        env[MODEL_NAME_ENV] = args.model
    env[CONFIG_NAME_ENV] = args.config_name
    env.setdefault("MINI_TRANSFORMER_UI_HOST", args.host)
    env.setdefault("MINI_TRANSFORMER_UI_PORT", str(args.port))

    app_path = Path(server_module.__file__).with_name("chainlit_app.py")
    cmd = [
        "chainlit",
        "run",
        str(app_path),
        "--host",
        args.host,
        "--port",
        str(args.port),
    ]
    try:
        proc = subprocess.run(cmd, env=env)
    except FileNotFoundError as exc:  # pragma: no cover
        raise SystemExit(
            "The 'chainlit' executable is missing. Install mini-transformer[server] to run the UI."
        ) from exc
    return proc.returncode


def fetch_main(argv: Sequence[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Download a model from the Hugging Face Hub")
    parser.add_argument(
        "model_id", help="Repo ID on Hugging Face Hub (e.g. AlaBoussoffara/transformer_small)"
    )
    parser.add_argument(
        "-n",
        "--name",
        help="Local directory name under trained_models/ (default derived from model id)",
    )
    parser.add_argument("--revision", help="Optional revision or branch to download")
    parser.add_argument("--repo-type", default="model", help="Hub repo type (default: model)")
    parser.add_argument("--token", help="Authentication token for private repos")
    parser.add_argument("--cache-dir", help="Custom Hugging Face cache directory")
    parser.add_argument("--model-root", help="Override path to trained_models root")
    parser.add_argument("--force", action="store_true", help="Overwrite existing local directory")
    args = parser.parse_args(argv)

    if args.model_root:
        os.environ[MODELS_ENV] = args.model_root

    try:
        destination = snapshot_to_local(
            args.model_id,
            local_name=args.name,
            revision=args.revision,
            repo_type=args.repo_type,
            token=args.token,
            cache_dir=args.cache_dir,
            force=args.force,
        )
    except FileExistsError as exc:  # pragma: no cover
        raise SystemExit(str(exc)) from exc

    print(f"Downloaded {args.model_id} → {destination}")
    config_path = destination / "configs" / f"{DEFAULT_CONFIG_NAME}.yaml"
    if not config_path.exists():
        print(
            f"[warning] Expected config file {config_path} not found. The inference CLI requires this file.",
            file=sys.stderr,
        )
    return 0


if __name__ == "__main__":  # pragma: no cover
    sys.exit(infer_main())