download_model.py

"""Descarga el repositorio GGUF de Qwen para uso local."""

from __future__ import annotations

import argparse
import os
from pathlib import Path

from huggingface_hub import snapshot_download


DEST_REPO = os.getenv(
    "LOCAL_MODEL_REPO",
    "AlbertiTechnology/qwen3-4b-instruct-gguf",
)
MODELS_DIR = Path(os.getenv("MODELS_DIR", "MODELS"))
LOCAL_MODEL_DIR = MODELS_DIR / "qwen3-4b-instruct-gguf"


def get_token(required: bool = False) -> str | None:
    token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN")
    if required and not token:
        raise RuntimeError("HF_TOKEN is required to create and upload the model repository.")
    return token


def download_local_model(
    repo_id: str = DEST_REPO,
    local_dir: Path = LOCAL_MODEL_DIR,
) -> Path:
    """Descarga el repositorio propio dentro de MODELS para inferencia local."""
    local_dir.parent.mkdir(parents=True, exist_ok=True)
    local_path = snapshot_download(
        repo_id=repo_id,
        token=get_token(),
        local_dir=local_dir,
    )
    print(f"Modelo local descargado en: {local_path}", flush=True)
    return Path(local_path)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--download-only",
        action="store_true",
        help="Sólo descarga DEST_REPO dentro de MODELS.",
    )
    parser.add_argument(
        "--repo",
        default=DEST_REPO,
        help="Repositorio que se descargara dentro de MODELS.",
    )
    parser.add_argument(
        "--local-dir",
        type=Path,
        default=LOCAL_MODEL_DIR,
        help="Directorio local donde se guardara el repositorio.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.download_only:
        download_local_model(args.repo, args.local_dir)
        return
    raise SystemExit("Use --download-only.")


if __name__ == "__main__":
    main()