311 lines
10 KiB
Python
311 lines
10 KiB
Python
"""README generation for quantised models.
|
|
|
|
Coordinates README creation by combining templates, formatting, and
|
|
original model information.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import re
|
|
from typing import TYPE_CHECKING
|
|
|
|
from helpers.logger import logger
|
|
from helpers.models.quantisation import QuantisationType
|
|
from helpers.readme.formatter import (
|
|
FileSizeFormatter,
|
|
TableFormatter,
|
|
TagFormatter,
|
|
)
|
|
from helpers.readme.templates import (
|
|
get_f16_row_template,
|
|
get_frontmatter_template,
|
|
get_header_template,
|
|
get_original_model_section,
|
|
get_quantisation_info,
|
|
)
|
|
from helpers.utils.config_parser import ConfigParser
|
|
|
|
if TYPE_CHECKING:
|
|
from pathlib import Path
|
|
|
|
from helpers.models.quantisation import ModelSource, QuantisationResult
|
|
|
|
# File size constant
|
|
GIBIBYTE = 1024**3
|
|
|
|
|
|
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, models_dir
|
|
)
|
|
|
|
readme_path.write_text(readme_content)
|
|
return readme_path
|
|
|
|
def _get_architecture(self, model_dir: Path) -> str | None:
|
|
"""Get the architecture from the model's config.json.
|
|
|
|
Returns:
|
|
Architecture name or None if not found.
|
|
"""
|
|
config_path = model_dir / "config.json"
|
|
if not config_path.exists():
|
|
return None
|
|
|
|
try:
|
|
with config_path.open(encoding="utf-8") as f:
|
|
config = json.load(f)
|
|
|
|
# Get the architectures field - it's a list
|
|
architectures = config.get("architectures", [])
|
|
if architectures:
|
|
arch_name = architectures[0]
|
|
# Get the mapped architecture (what it will be converted to)
|
|
parser = ConfigParser()
|
|
mapped_arch = parser.get_architecture_mapping(arch_name)
|
|
logger.info(f"Architecture: {arch_name} -> {mapped_arch}")
|
|
return mapped_arch
|
|
|
|
except Exception as e:
|
|
logger.warning(f"Could not determine architecture: {e}")
|
|
|
|
return None
|
|
|
|
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)
|
|
content["readme"] = readme_content
|
|
logger.info(f"Preserved original README ({len(readme_content)} characters)")
|
|
else:
|
|
# This is our README, try to extract original content
|
|
logger.info("Found existing generated README, extracting original content")
|
|
# Try to find the separator
|
|
separator_idx = readme_content.find("\n---\n\n## Original Model Information\n")
|
|
if separator_idx > 0:
|
|
content["readme"] = readme_content[separator_idx + 37 :]
|
|
else:
|
|
logger.info("No README found to preserve")
|
|
|
|
# Parse frontmatter if we have content
|
|
if content["readme"]:
|
|
parsed = self._parse_frontmatter(content["readme"])
|
|
content.update(parsed)
|
|
|
|
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,
|
|
models_dir: Path | 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
|
|
tag_formatter = TagFormatter()
|
|
original_tags = original_content["tags"].split(",") if original_content["tags"] else []
|
|
all_tags = tag_formatter.build_tags(results, original_tags)
|
|
|
|
# Build frontmatter
|
|
content = get_frontmatter_template(
|
|
original_content["licence"],
|
|
model_source.source_model,
|
|
all_tags,
|
|
)
|
|
|
|
# Add header
|
|
content += get_header_template(
|
|
model_source.original_author,
|
|
model_source.model_name,
|
|
model_source.source_model,
|
|
)
|
|
|
|
# Add quantisation table
|
|
table_formatter = TableFormatter()
|
|
for quant_type in table_formatter.get_ordered_quantisation_types():
|
|
result = results.get(quant_type)
|
|
content += table_formatter.format_quantisation_row(
|
|
quant_type, result, model_source, output_repo
|
|
)
|
|
|
|
# Add F16 row if applicable
|
|
if not model_source.is_gguf_repo and output_repo:
|
|
content += self._format_f16_row(model_source, results, output_repo, models_dir)
|
|
|
|
# Add quantisation information
|
|
content += get_quantisation_info()
|
|
|
|
# Add original model section if available
|
|
if original_content.get("readme"):
|
|
content += get_original_model_section(original_content["readme"])
|
|
|
|
return content
|
|
|
|
def _format_f16_row(
|
|
self,
|
|
model_source: ModelSource,
|
|
results: dict[QuantisationType, QuantisationResult],
|
|
output_repo: str,
|
|
models_dir: Path | None = None,
|
|
) -> str:
|
|
"""Format F16 GGUF row for the table.
|
|
|
|
Creates a properly formatted F16 reference row for the quantisation
|
|
table using source model information, results data, and repository
|
|
details with optional models directory for file size calculation.
|
|
|
|
Returns:
|
|
Formatted F16 table row.
|
|
"""
|
|
# Get F16 result from results dict
|
|
f16_result = results.get(QuantisationType.F16)
|
|
|
|
# Get file size
|
|
f16_size = "-"
|
|
if f16_result and hasattr(f16_result, "file_size"):
|
|
f16_size = f16_result.file_size or "-"
|
|
elif models_dir:
|
|
# Try to get from actual file
|
|
f16_filename = f"{model_source.original_author}-{model_source.model_name}-f16.gguf"
|
|
f16_path = models_dir / model_source.model_name / f16_filename
|
|
if f16_path.exists():
|
|
f16_size = FileSizeFormatter.get_file_size(f16_path)
|
|
|
|
# Get status
|
|
status = "planned"
|
|
if f16_result and hasattr(f16_result, "status"):
|
|
status = f16_result.status
|
|
|
|
return get_f16_row_template(
|
|
model_source.original_author,
|
|
model_source.model_name,
|
|
output_repo,
|
|
f16_size,
|
|
status,
|
|
)
|