soup/soup_cli/commands/push.py

"""soup push — upload a trained model to HuggingFace Hub."""

import json
import os
from pathlib import Path
from typing import Optional

import typer
from rich.console import Console
from rich.panel import Panel

console = Console()

# Files that should exist in a valid LoRA adapter directory
ADAPTER_FILES = {"adapter_config.json", "adapter_model.safetensors"}
ADAPTER_FILES_ALT = {"adapter_config.json", "adapter_model.bin"}


def push(
    model: str = typer.Option(
        ...,
        "--model",
        "-m",
        help="Path to the trained model / LoRA adapter directory",
    ),
    repo: str = typer.Option(
        ...,
        "--repo",
        "-r",
        help="HuggingFace repo ID, e.g. username/my-model",
    ),
    private: bool = typer.Option(
        False,
        "--private",
        help="Make the HuggingFace repo private",
    ),
    token: Optional[str] = typer.Option(
        None,
        "--token",
        "-t",
        help="[deprecated] Use HF_TOKEN env var instead. Falls back to cached login.",
        envvar="HF_TOKEN",
    ),
    commit_message: str = typer.Option(
        "Upload model trained with Soup CLI",
        "--message",
        help="Commit message for the upload",
    ),
):
    """Push a trained model to HuggingFace Hub."""
    model_path = Path(model)

    # --- Validate model directory ---
    if not model_path.exists():
        console.print(f"[red]Model path not found: {model_path}[/]")
        raise typer.Exit(1)

    if not model_path.is_dir():
        console.print(f"[red]Expected a directory, got a file: {model_path}[/]")
        raise typer.Exit(1)

    files_in_dir = {f.name for f in model_path.iterdir() if f.is_file()}
    is_adapter = ADAPTER_FILES.issubset(files_in_dir) or ADAPTER_FILES_ALT.issubset(files_in_dir)

    if not is_adapter and "config.json" not in files_in_dir:
        console.print(
            "[red]Directory does not look like a valid model or LoRA adapter.[/]\n"
            "Expected adapter_config.json (LoRA) or config.json (full model)."
        )
        raise typer.Exit(1)

    # --- Resolve HF token ---
    hf_token = token or os.environ.get("HF_TOKEN")
    if not hf_token:
        hf_token = _get_cached_token()

    if not hf_token:
        console.print(
            "[red]No HuggingFace token found.[/]\n"
            "Provide one via:\n"
            "  --token YOUR_TOKEN\n"
            "  HF_TOKEN=... env variable\n"
            "  huggingface-cli login"
        )
        raise typer.Exit(1)

    # --- Show upload plan ---
    file_count = sum(1 for _ in model_path.rglob("*") if _.is_file())
    total_size = sum(f.stat().st_size for f in model_path.rglob("*") if f.is_file())
    size_str = _format_size(total_size)

    console.print(
        Panel(
            f"Source:  [bold]{model_path}[/]\n"
            f"Repo:   [bold]{repo}[/]\n"
            f"Type:   [bold]{'LoRA adapter' if is_adapter else 'Full model'}[/]\n"
            f"Files:  [bold]{file_count}[/]\n"
            f"Size:   [bold]{size_str}[/]\n"
            f"Private: [bold]{private}[/]",
            title="Upload Plan",
        )
    )

    # --- Upload ---
    console.print("[dim]Uploading to HuggingFace Hub...[/]")

    try:
        from huggingface_hub import HfApi

        api = HfApi(token=hf_token)

        # Create repo if it doesn't exist
        api.create_repo(repo_id=repo, private=private, exist_ok=True)

        # Upload the entire directory
        api.upload_folder(
            folder_path=str(model_path),
            repo_id=repo,
            commit_message=commit_message,
        )

        # Generate and upload model card if not present
        readme_path = model_path / "README.md"
        if not readme_path.exists():
            model_card = _generate_model_card(model_path, repo, is_adapter)
            api.upload_file(
                path_or_fileobj=model_card.encode("utf-8"),
                path_in_repo="README.md",
                repo_id=repo,
                commit_message="Add model card (generated by Soup CLI)",
            )

    except ImportError:
        console.print(
            "[red]huggingface-hub not installed.[/]\n"
            "Run: [bold]pip install huggingface-hub[/]"
        )
        raise typer.Exit(1)
    except Exception as exc:
        console.print(f"[red]Upload failed: {exc}[/]")
        raise typer.Exit(1)

    repo_url = f"https://huggingface.co/{repo}"
    console.print(
        Panel(
            f"Repo: [bold blue]{repo_url}[/]\n\n"
            f"Use it:\n"
            f"  [bold]soup chat --model {repo}[/]\n"
            f"  [bold]from peft import PeftModel[/]",
            title="[bold green]Upload Complete![/]",
        )
    )


def _get_cached_token() -> Optional[str]:
    """Try to read HF token from cached login."""
    token_path = Path.home() / ".huggingface" / "token"
    if token_path.exists():
        return token_path.read_text().strip()

    # New location used by huggingface_hub
    token_path_new = Path.home() / ".cache" / "huggingface" / "token"
    if token_path_new.exists():
        return token_path_new.read_text().strip()

    return None


def _format_size(size_bytes: int) -> str:
    """Format bytes into human-readable string."""
    for unit in ("B", "KB", "MB", "GB"):
        if size_bytes < 1024:
            return f"{size_bytes:.1f} {unit}"
        size_bytes /= 1024
    return f"{size_bytes:.1f} TB"


def _generate_model_card(model_path: Path, repo_id: str, is_adapter: bool) -> str:
    """Generate a basic model card README."""
    adapter_info = ""
    config_path = model_path / "adapter_config.json"
    if is_adapter and config_path.exists():
        try:
            with open(config_path, encoding="utf-8") as f:
                config = json.load(f)
            base = config.get("base_model_name_or_path", "unknown")
            lora_r = config.get("r", "?")
            lora_alpha = config.get("lora_alpha", "?")
            adapter_info = (
                f"- **Base model:** `{base}`\n"
                f"- **LoRA rank:** {lora_r}\n"
                f"- **LoRA alpha:** {lora_alpha}\n"
            )
        except (json.JSONDecodeError, OSError):
            pass

    model_name = repo_id.split("/")[-1] if "/" in repo_id else repo_id

    return f"""---
tags:
  - soup-cli
  - fine-tuned
  - lora
library_name: peft
---

# {model_name}

Fine-tuned model uploaded with [Soup CLI](https://github.com/MakazhanAlpamys/Soup).

## Model Details

{adapter_info if adapter_info else "This is a fine-tuned language model."}

## Usage

```python
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("BASE_MODEL")
model = PeftModel.from_pretrained(model, "{repo_id}")
tokenizer = AutoTokenizer.from_pretrained("{repo_id}")
```

Or with Soup CLI:

```bash
soup chat --model {repo_id}
```

## Training

Trained using [Soup CLI](https://github.com/MakazhanAlpamys/Soup) — fine-tune LLMs in one command.
"""