613 lines
21 KiB
Python
613 lines
21 KiB
Python
"""HuggingFace operations service.
|
|
|
|
Handles all interactions with HuggingFace including model downloads,
|
|
uploads, README generation, and repository management. Uses UK English
|
|
spelling conventions throughout.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
import shutil
|
|
import subprocess
|
|
import tempfile
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING
|
|
|
|
from helpers.config.quantisation_configs import QUANTISATION_CONFIGS
|
|
from helpers.logger import logger
|
|
from helpers.models.quantisation import QuantisationType
|
|
|
|
if TYPE_CHECKING:
|
|
from helpers.models.quantisation import ModelSource, QuantisationResult
|
|
|
|
# Constants for file size formatting
|
|
GIBIBYTE = 1024**3
|
|
|
|
|
|
class HuggingFaceService:
|
|
"""Manages HuggingFace repository operations.
|
|
|
|
Provides methods for downloading models, uploading files, and managing
|
|
repositories. Handles authentication, error recovery, and progress tracking
|
|
for robust interaction with HuggingFace services.
|
|
"""
|
|
|
|
@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.
|
|
"""
|
|
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 upload_file(
|
|
repo_id: str,
|
|
local_path: Path,
|
|
repo_path: str | None = None,
|
|
create_repo: bool = False,
|
|
) -> None:
|
|
"""Upload a file to HuggingFace repository.
|
|
|
|
Uploads a single file to the specified repository path. Can create
|
|
the repository if it doesn't exist. Uses git directly when possible
|
|
to avoid automatic PR creation.
|
|
|
|
Raises:
|
|
CalledProcessError: If upload fails.
|
|
"""
|
|
repo_path = repo_path or local_path.name
|
|
logger.info(f"Uploading {local_path.name} to {repo_id}/{repo_path}")
|
|
|
|
# Try git-based upload first to avoid PR creation
|
|
if HuggingFaceService._try_git_upload(
|
|
repo_id, local_path, repo_path, create_repo=create_repo
|
|
):
|
|
logger.info(f"Uploaded {repo_path} via git")
|
|
return
|
|
|
|
# Fallback to huggingface-cli
|
|
logger.info("Git upload failed, trying huggingface-cli...")
|
|
cmd = [
|
|
"huggingface-cli",
|
|
"upload",
|
|
repo_id,
|
|
str(local_path),
|
|
repo_path,
|
|
"--revision",
|
|
"main", # Explicitly push to main branch
|
|
"--commit-message",
|
|
f"Add {repo_path}",
|
|
]
|
|
|
|
if create_repo:
|
|
cmd.append("--create")
|
|
|
|
try:
|
|
subprocess.run(cmd, check=True, capture_output=True)
|
|
logger.info(f"Uploaded {repo_path}")
|
|
except subprocess.CalledProcessError:
|
|
if create_repo:
|
|
# Repository might already exist, retry without --create
|
|
cmd = cmd[:-1] # Remove --create flag
|
|
subprocess.run(cmd, check=True, capture_output=True, text=True)
|
|
logger.info(f"Updated {repo_path}")
|
|
else:
|
|
raise
|
|
|
|
@staticmethod
|
|
def _try_git_upload(
|
|
repo_id: str,
|
|
local_path: Path,
|
|
repo_path: str,
|
|
*,
|
|
create_repo: bool = False,
|
|
) -> bool:
|
|
"""Try to upload file using git directly to avoid PR creation.
|
|
|
|
Returns:
|
|
bool: True if upload successful, False if should fallback to CLI.
|
|
"""
|
|
try:
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
temp_path = Path(temp_dir)
|
|
repo_url = f"https://huggingface.co/{repo_id}"
|
|
|
|
# Clone repository
|
|
logger.info(f"Cloning {repo_url}...")
|
|
result = subprocess.run(
|
|
["git", "clone", repo_url, str(temp_path / "repo")],
|
|
check=False,
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
|
|
if result.returncode != 0:
|
|
if create_repo:
|
|
# Repository doesn't exist, let huggingface-cli handle creation
|
|
return False
|
|
logger.warning(f"Clone failed: {result.stderr}")
|
|
return False
|
|
|
|
repo_dir = temp_path / "repo"
|
|
target_file = repo_dir / repo_path
|
|
|
|
# Ensure target directory exists
|
|
target_file.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Copy file
|
|
shutil.copy2(local_path, target_file)
|
|
|
|
# Check if there are any changes
|
|
status_result = subprocess.run(
|
|
["git", "status", "--porcelain"],
|
|
cwd=repo_dir,
|
|
capture_output=True,
|
|
text=True,
|
|
check=True,
|
|
)
|
|
|
|
if not status_result.stdout.strip():
|
|
logger.info(f"No changes detected for {repo_path}, file already up-to-date")
|
|
return True # File is already up-to-date, no need to push
|
|
|
|
# Git add, commit, push
|
|
subprocess.run(
|
|
["git", "add", repo_path],
|
|
cwd=repo_dir,
|
|
check=True,
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
subprocess.run(
|
|
["git", "commit", "-m", f"Update {repo_path}"],
|
|
cwd=repo_dir,
|
|
check=True,
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
subprocess.run(
|
|
["git", "push"],
|
|
cwd=repo_dir,
|
|
check=True,
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
|
|
return True
|
|
|
|
except subprocess.CalledProcessError as e:
|
|
logger.warning(f"Git upload failed: {e}")
|
|
return False
|
|
except Exception as e:
|
|
logger.warning(f"Git upload error: {e}")
|
|
return False
|
|
|
|
|
|
class ReadmeGenerator:
|
|
"""Generates README files for quantised models.
|
|
|
|
Creates comprehensive README documentation including model cards,
|
|
quantisation details, and status tracking. Supports both initial
|
|
planning documentation and final result summaries.
|
|
"""
|
|
|
|
def generate(
|
|
self,
|
|
model_source: ModelSource,
|
|
results: dict[QuantisationType, QuantisationResult],
|
|
models_dir: Path,
|
|
output_repo: str | None = None,
|
|
) -> Path:
|
|
"""Generate README file for quantised model repository.
|
|
|
|
Creates a comprehensive README with frontmatter, quantisation table,
|
|
and original model information. Handles status tracking for planned,
|
|
processing, and completed quantisations.
|
|
|
|
Returns:
|
|
Path to generated README file.
|
|
"""
|
|
logger.info("Creating model card...")
|
|
|
|
model_dir = models_dir / model_source.model_name
|
|
readme_path = model_dir / "README.md"
|
|
|
|
# Get original README content
|
|
original_content = self._get_original_readme(model_source, model_dir)
|
|
|
|
# Generate new README
|
|
readme_content = self._generate_readme_content(
|
|
model_source, results, original_content, output_repo
|
|
)
|
|
|
|
readme_path.write_text(readme_content)
|
|
return readme_path
|
|
|
|
def _get_original_readme(self, model_source: ModelSource, model_dir: Path) -> dict[str, str]:
|
|
"""Extract original README and metadata.
|
|
|
|
Downloads or reads the original model's README for inclusion in the
|
|
quantised model documentation. Parses YAML frontmatter if present.
|
|
|
|
Returns:
|
|
Dictionary with readme content, licence, tags, and frontmatter.
|
|
"""
|
|
content = {"readme": "", "licence": "apache-2.0", "tags": "", "frontmatter": ""}
|
|
|
|
# Check for preserved original README first
|
|
original_readme_path = model_dir / "README.original.md"
|
|
readme_path = model_dir / "README.md"
|
|
|
|
if original_readme_path.exists():
|
|
# Use the preserved original
|
|
content["readme"] = original_readme_path.read_text(encoding="utf-8")
|
|
logger.info(f"Found preserved original README ({len(content['readme'])} characters)")
|
|
elif readme_path.exists():
|
|
# First time - preserve the original and use it
|
|
readme_content = readme_path.read_text(encoding="utf-8")
|
|
|
|
# Check if this is already our generated README
|
|
if (
|
|
f"{model_source.original_author}-{model_source.model_name}-GGUF"
|
|
not in readme_content
|
|
):
|
|
# This is the original - preserve it
|
|
original_readme_path.write_text(readme_content, encoding="utf-8")
|
|
content["readme"] = readme_content
|
|
readme_len = len(content["readme"])
|
|
logger.info(
|
|
f"Preserved original README as README.original.md ({readme_len} characters)"
|
|
)
|
|
else:
|
|
# This is our generated README, need to download the original
|
|
logger.info("Found generated README, downloading original from source")
|
|
content = self._download_readme(model_source)
|
|
# Save the downloaded original for future use
|
|
if content["readme"]:
|
|
original_readme_path.write_text(content["readme"], encoding="utf-8")
|
|
logger.info("Preserved downloaded original README as README.original.md")
|
|
else:
|
|
# No local README - download from source
|
|
content = self._download_readme(model_source)
|
|
# Save the downloaded original for future use
|
|
if content["readme"]:
|
|
original_readme_path.write_text(content["readme"], encoding="utf-8")
|
|
logger.info("Preserved downloaded original README as README.original.md")
|
|
|
|
# Parse frontmatter if present
|
|
if content["readme"].startswith("---\n"):
|
|
content = self._parse_frontmatter(content["readme"])
|
|
|
|
return content
|
|
|
|
def _download_readme(self, model_source: ModelSource) -> dict[str, str]:
|
|
"""Download README from HuggingFace repository.
|
|
|
|
Attempts to download just the README.md file from the source repository
|
|
for efficient documentation extraction.
|
|
|
|
Returns:
|
|
Dictionary with readme content and default metadata.
|
|
"""
|
|
content = {"readme": "", "licence": "apache-2.0", "tags": "", "frontmatter": ""}
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
try:
|
|
logger.info(f"Downloading README from {model_source.source_model}...")
|
|
subprocess.run(
|
|
[
|
|
"huggingface-cli",
|
|
"download",
|
|
model_source.source_model,
|
|
"--include",
|
|
"README.md",
|
|
"--local-dir",
|
|
temp_dir,
|
|
],
|
|
check=True,
|
|
capture_output=True,
|
|
)
|
|
|
|
readme_path = Path(temp_dir) / "README.md"
|
|
if readme_path.exists():
|
|
content["readme"] = readme_path.read_text(encoding="utf-8")
|
|
logger.info(f"Downloaded README ({len(content['readme'])} characters)")
|
|
except subprocess.CalledProcessError as e:
|
|
logger.warning(f"Failed to download README: {e}")
|
|
|
|
return content
|
|
|
|
def _parse_frontmatter(self, readme_text: str) -> dict[str, str]:
|
|
"""Parse YAML frontmatter from README.
|
|
|
|
Extracts metadata from YAML frontmatter including licence, tags,
|
|
and other model card fields.
|
|
|
|
Returns:
|
|
Dictionary with separated content and metadata.
|
|
"""
|
|
lines = readme_text.split("\n")
|
|
if lines[0] != "---":
|
|
return {
|
|
"readme": readme_text,
|
|
"licence": "apache-2.0",
|
|
"tags": "",
|
|
"frontmatter": "",
|
|
}
|
|
|
|
frontmatter_end = -1
|
|
for i, line in enumerate(lines[1:], 1):
|
|
if line == "---":
|
|
frontmatter_end = i
|
|
break
|
|
|
|
if frontmatter_end == -1:
|
|
return {
|
|
"readme": readme_text,
|
|
"licence": "apache-2.0",
|
|
"tags": "",
|
|
"frontmatter": "",
|
|
}
|
|
|
|
frontmatter = "\n".join(lines[1:frontmatter_end])
|
|
content = "\n".join(lines[frontmatter_end + 1 :])
|
|
|
|
# Extract licence
|
|
licence_match = re.search(r"^license:\s*(.+)$", frontmatter, re.MULTILINE)
|
|
licence_val = licence_match.group(1).strip().strip('"') if licence_match else "apache-2.0"
|
|
|
|
# Extract tags
|
|
tags = []
|
|
in_tags = False
|
|
for line in frontmatter.split("\n"):
|
|
if line.startswith("tags:"):
|
|
in_tags = True
|
|
continue
|
|
if in_tags:
|
|
if line.startswith("- "):
|
|
tags.append(line[2:].strip())
|
|
elif line and not line.startswith(" "):
|
|
break
|
|
|
|
return {
|
|
"readme": content,
|
|
"licence": licence_val,
|
|
"tags": ",".join(tags),
|
|
"frontmatter": frontmatter,
|
|
}
|
|
|
|
def _generate_readme_content(
|
|
self,
|
|
model_source: ModelSource,
|
|
results: dict[QuantisationType, QuantisationResult],
|
|
original_content: dict[str, str],
|
|
output_repo: str | None = None,
|
|
) -> str:
|
|
"""Generate complete README content with quantisation details.
|
|
|
|
Creates the full README including YAML frontmatter, quantisation status
|
|
table, and original model information.
|
|
|
|
Returns:
|
|
Complete README markdown content.
|
|
"""
|
|
# Build tags
|
|
our_tags = [
|
|
"quantised",
|
|
"gguf",
|
|
"q3_k_m",
|
|
"q3_k_l",
|
|
"q3_k_xl",
|
|
"q4_k_m",
|
|
"q4_k_l",
|
|
"q5_k_m",
|
|
"q5_k_l",
|
|
"q6_k",
|
|
"q6_k_l",
|
|
"q8_0",
|
|
"bartowski-method",
|
|
]
|
|
original_tags = original_content["tags"].split(",") if original_content["tags"] else []
|
|
all_tags = sorted(set(our_tags + original_tags))
|
|
|
|
# Build frontmatter
|
|
frontmatter = f"""---
|
|
license: {original_content["licence"]}
|
|
library_name: gguf
|
|
base_model: {model_source.source_model}
|
|
tags:
|
|
"""
|
|
for tag in all_tags:
|
|
if tag.strip():
|
|
frontmatter += f"- {tag.strip()}\n"
|
|
|
|
frontmatter += "---\n\n"
|
|
|
|
# Build main content
|
|
hf_url = f"https://huggingface.co/{model_source.source_model}"
|
|
content = f"""# {model_source.original_author}-{model_source.model_name}-GGUF
|
|
|
|
GGUF quantisations of [{model_source.source_model}]({hf_url}) using
|
|
[Bartowski](https://huggingface.co/bartowski)'s method. Created with [llm-gguf-tools](https://git.tomfos.tr/tom/llm-gguf-tools)
|
|
which replicates Bartowski's quantisation profiles.
|
|
|
|
| Variant | Configuration | File Size | Status |
|
|
|---|---|---|---|
|
|
"""
|
|
|
|
# Add results table - group by layer config patterns
|
|
supported_types = [
|
|
QuantisationType.Q3_K_M,
|
|
QuantisationType.Q3_K_L,
|
|
QuantisationType.Q3_K_XL,
|
|
QuantisationType.Q4_K_M,
|
|
QuantisationType.Q4_K_L,
|
|
QuantisationType.Q5_K_M,
|
|
QuantisationType.Q5_K_L,
|
|
QuantisationType.Q6_K,
|
|
QuantisationType.Q6_K_L,
|
|
QuantisationType.Q8_0,
|
|
]
|
|
|
|
for quant_type in supported_types:
|
|
result = results.get(quant_type)
|
|
if not result:
|
|
result = type("Result", (), {"status": "planned", "success": False})()
|
|
|
|
config = QUANTISATION_CONFIGS.get(quant_type)
|
|
file_size = self._format_file_size(result)
|
|
status = self._format_status(result, model_source, quant_type, output_repo)
|
|
|
|
# Get configuration description from the config itself
|
|
config_desc = config.get_compact_config(QUANTISATION_CONFIGS) if config else f"{quant_type} all layers"
|
|
|
|
content += f"| **{quant_type.value}** | {config_desc} | {file_size} | {status} |\n"
|
|
|
|
content += """
|
|
|
|
**Key:** `E` = Embeddings, `O` = Output, `A` = Attention, `F` = FFN
|
|
|
|
See [Bartowski Analysis](https://git.tomfos.tr/tom/llm-gguf-tools/src/branch/main/docs/bartowski_analysis.md)
|
|
for detailed quantisation strategies and [Documentation](https://git.tomfos.tr/tom/llm-gguf-tools/src/branch/main/docs/)
|
|
for more on the tools and methods I use.
|
|
|
|
"""
|
|
|
|
# Add original content
|
|
if original_content["readme"]:
|
|
content += "## Original Model Card\n\n---\n\n" + original_content["readme"]
|
|
else:
|
|
content += f"## Original Model\n\nQuantisation of [{model_source.source_model}](https://huggingface.co/{model_source.source_model})."
|
|
|
|
return frontmatter + content
|
|
|
|
def _format_file_size(self, result: QuantisationResult) -> str:
|
|
"""Format file size for README table.
|
|
|
|
Returns:
|
|
Formatted file size string or dash if not available.
|
|
"""
|
|
if hasattr(result, "file_size") and result.file_size:
|
|
return result.file_size
|
|
if hasattr(result, "success") and result.success and hasattr(result, "file_path"):
|
|
# Try to get file size from path if available
|
|
try:
|
|
if result.file_path and Path(result.file_path).exists():
|
|
size_bytes = Path(result.file_path).stat().st_size
|
|
size_gb = size_bytes / GIBIBYTE
|
|
return f"{size_gb:.1f}GB"
|
|
except Exception:
|
|
pass
|
|
return "-"
|
|
|
|
def _format_status(
|
|
self,
|
|
result: QuantisationResult,
|
|
model_source: ModelSource,
|
|
quant_type: QuantisationType,
|
|
output_repo: str | None,
|
|
) -> str:
|
|
"""Format status indicator for README table.
|
|
|
|
Creates appropriate status indicator based on quantisation state
|
|
including progress indicators, file sizes, and download links.
|
|
|
|
Returns:
|
|
Formatted status string for table cell.
|
|
"""
|
|
status_map = {
|
|
"planned": "⏳ Queued",
|
|
"processing": "🔄 Processing...",
|
|
"uploading": "⬆️ Uploading...",
|
|
"failed": "❌ Failed",
|
|
}
|
|
|
|
if hasattr(result, "status") and result.status in status_map:
|
|
base_status = status_map[result.status]
|
|
|
|
if result.status == "uploading" and hasattr(result, "file_size") and result.file_size:
|
|
return f"{base_status} ({result.file_size})"
|
|
if result.status == "completed" or (hasattr(result, "success") and result.success):
|
|
return self._format_success_status(result, model_source, quant_type, output_repo)
|
|
return base_status
|
|
|
|
# Legacy support
|
|
if hasattr(result, "success") and result.success:
|
|
return self._format_success_status(result, model_source, quant_type, output_repo)
|
|
return "❌ Failed"
|
|
|
|
def _format_success_status(
|
|
self,
|
|
result: QuantisationResult,
|
|
model_source: ModelSource,
|
|
quant_type: QuantisationType,
|
|
output_repo: str | None,
|
|
) -> str:
|
|
"""Format successful quantisation status with download link.
|
|
|
|
Creates a download link if repository information is available,
|
|
otherwise shows file size.
|
|
|
|
Returns:
|
|
Formatted success status string.
|
|
"""
|
|
if not output_repo:
|
|
return (
|
|
f"✅ {result.file_size}"
|
|
if hasattr(result, "file_size") and result.file_size
|
|
else "✅ Available"
|
|
)
|
|
|
|
filename = (
|
|
f"{model_source.original_author}-{model_source.model_name}-{quant_type.value}.gguf"
|
|
)
|
|
url = f"https://huggingface.co/{output_repo}?show_file_info={filename}"
|
|
|
|
if hasattr(result, "file_size") and result.file_size:
|
|
return f"[✅ {result.file_size}]({url})"
|
|
return f"[✅ Available]({url})"
|