124 lines
3.8 KiB
Python
124 lines
3.8 KiB
Python
"""HuggingFace API client operations.
|
|
|
|
Provides basic HuggingFace API operations including authentication,
|
|
model downloads, and user information retrieval.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import subprocess
|
|
from typing import TYPE_CHECKING
|
|
|
|
from helpers.logger import logger
|
|
|
|
if TYPE_CHECKING:
|
|
from pathlib import Path
|
|
|
|
|
|
class HuggingFaceClient:
|
|
"""Manages basic HuggingFace API operations.
|
|
|
|
Provides methods for authentication verification, model downloads,
|
|
and user information retrieval using the HuggingFace CLI.
|
|
"""
|
|
|
|
@staticmethod
|
|
def get_username() -> str:
|
|
"""Get authenticated HuggingFace username.
|
|
|
|
Retrieves the current user's HuggingFace username using the CLI.
|
|
Requires prior authentication via `huggingface-cli login`.
|
|
|
|
Returns:
|
|
HuggingFace username.
|
|
|
|
Raises:
|
|
RuntimeError: If not authenticated or CLI not available.
|
|
"""
|
|
try:
|
|
result = subprocess.run(
|
|
["huggingface-cli", "whoami"],
|
|
capture_output=True,
|
|
text=True,
|
|
check=True,
|
|
)
|
|
return result.stdout.strip()
|
|
except (subprocess.CalledProcessError, FileNotFoundError) as err:
|
|
msg = "Please log in to HuggingFace first: huggingface-cli login"
|
|
raise RuntimeError(msg) from err
|
|
|
|
@staticmethod
|
|
def download_model(
|
|
model_name: str,
|
|
output_dir: Path,
|
|
include_pattern: str | None = None,
|
|
) -> None:
|
|
"""Download model from HuggingFace.
|
|
|
|
Downloads a complete model or specific files matching a pattern.
|
|
Creates the output directory if it doesn't exist. Supports filtered
|
|
downloads for efficient bandwidth usage when only certain files are needed.
|
|
The model identifier follows HuggingFace naming conventions (e.g. "meta-llama/Llama-2-7b").
|
|
"""
|
|
logger.info(f"Downloading {model_name} to {output_dir}")
|
|
|
|
cmd = [
|
|
"huggingface-cli",
|
|
"download",
|
|
model_name,
|
|
"--local-dir",
|
|
str(output_dir),
|
|
]
|
|
|
|
if include_pattern:
|
|
cmd.extend(["--include", include_pattern])
|
|
|
|
subprocess.run(cmd, check=True, capture_output=True, text=True)
|
|
logger.info("Download complete")
|
|
|
|
@staticmethod
|
|
def check_authentication() -> bool:
|
|
"""Check if user is authenticated with HuggingFace.
|
|
|
|
Returns:
|
|
True if authenticated, False otherwise.
|
|
"""
|
|
try:
|
|
result = subprocess.run(
|
|
["huggingface-cli", "whoami"],
|
|
capture_output=True,
|
|
text=True,
|
|
check=False,
|
|
)
|
|
except FileNotFoundError:
|
|
logger.error(
|
|
"huggingface-cli not found. Please install with: pip install huggingface-hub"
|
|
)
|
|
return False
|
|
else:
|
|
return result.returncode == 0
|
|
|
|
@staticmethod
|
|
def get_model_info(model_id: str) -> dict | None:
|
|
"""Get model information from HuggingFace.
|
|
|
|
Retrieves metadata about a model from the HuggingFace Hub using the
|
|
CLI interface. Returns the model information as a dictionary if found.
|
|
|
|
Returns:
|
|
Model information dictionary or None if not found.
|
|
"""
|
|
try:
|
|
# Use huggingface-cli to get model info
|
|
result = subprocess.run(
|
|
["huggingface-cli", "model-info", model_id],
|
|
capture_output=True,
|
|
text=True,
|
|
check=True,
|
|
)
|
|
except subprocess.CalledProcessError:
|
|
logger.warning(f"Could not get info for model: {model_id}")
|
|
return None
|
|
else:
|
|
# Parse the output (this is simplified - actual implementation would parse JSON)
|
|
return {"output": result.stdout}
|