dl_model.py

"""Download modles."""
# pylint: disable=invalid-name, broad-exception-caught, line-too-long
from typing import Optional

import typer
from dl_hf_model import dl_hf_model
from loguru import logger

url = "https://huggingface.co/TheBloke/Upstage-Llama-2-70B-instruct-v2-GGML/blob/main/upstage-llama-2-70b-instruct-v2.ggmlv3.q3_K_S.bin"


__version__ = "0.0.1"

app = typer.Typer(
    name="dl-mode",
    add_completion=False,
    help="donwload models from hf and save to a dir (default models)",
)


def _version_callback(value: bool) -> None:
    if value:
        typer.echo(
            f"{app.info.name} v.{__version__} -- download models for given url(s)"
        )
        raise typer.Exit()


@app.command()
def main(
    urls: str = typer.Argument(  # pylint: disable=unused-argument
        "",
        help=f"one or more urls (default {url})",
        show_default=False,
    ),
    version: Optional[bool] = typer.Option(  # pylint: disable=unused-argument
        None,
        "--version",
        "-v",
        "-V",
        help="Show version info and exit.",
        callback=_version_callback,
        is_eager=True,
    ),
    model_dir: Optional[str] = typer.Option(
        None,
        "--mode-dir",
        help="dir to save downloaded models (default models)",
    ),
):
    """Download a model or model given url(s)."""
    logger.trace(f"{urls}")
    if model_dir is None:
        model_dir = "models"
    if isinstance(urls, str):
        urls.split()

    url_list = urls[:]
    if not urls:
        url_list = [url]
    try:
        for elm in url_list:
            dl_hf_model(elm)
    except Exception as exc:
        logger.error(exc)
        raise typer.Exit()


if __name__ == "__main__":
    try:
        app()
    except Exception as exc_:
        logger.error(exc_)