"""Model Harvester: explore Hugging Face for MLX/GGUF models that genuinely fit this machine, deduplicate them and propose a download queue. Never downloads blindly.""" from __future__ import annotations import asyncio import json import logging import os import re import time from typing import Any from .config import Settings from .downloads import TRUSTED_GGUF_AUTHORS, pick_gguf_files from .jobs import Job, JobRunner from .models import compat, formats from .models.estimator import kv_bytes_per_token log = logging.getLogger("llm_api.harvester") GB = 1024**3 DEFAULT_SOURCES = { "mlx": ["mlx-community"], "gguf": ["unsloth", "bartowski", "ggml-org", "lmstudio-community"], } TASK_HINTS = { "coding": re.compile(r"coder|code|devstral|codestral|starcoder|deepseek-coder|kat-coder", re.I), "reasoning": re.compile(r"thinking|reason|r1|qwq|deepseek-r|magistral|-think", re.I), "vision": re.compile(r"-vl|vision|pixtral|llava|gemma-?3|gemma-?4|qwen3\.[5-8]|paligemma|-omni|ocr", re.I), "embedding": re.compile(r"embed|e5-|bge|gte-|minilm|nomic", re.I), "reranker": re.compile(r"rerank", re.I), "multilingual": re.compile(r"multilingual|qwen|gemma|aya|mistral", re.I), } EXCLUDE = re.compile(r"uncensored|abliterated|heretic|nsfw|erotic|roleplay|-rp-|waifu|distill(ed)?-mlx-4bit-claude|" r"whisper|parakeet|tts|kokoro|orpheus|-asr|speech|audio|text-to-image|qwen-image|flux|stable-diffusion|" r"sd-|wan2|video|roformer|-base(-|$)|pretrain|draft|eagle|mtp-", re.I) def _base_key(repo: str, tags: list[str]) -> str: bm = next((t.split(":", 2)[2] for t in tags if t.startswith("base_model:quantized:")), None) if not bm: bm = next((t.split(":", 1)[1] for t in tags if t.startswith("base_model:") and "finetune:" not in t), None) if bm: return bm.split("/")[-1].lower() name = repo.split("/")[-1].lower() name = re.sub(r"-(\d+bit|q\d[_a-z0-9]*|iq\d[_a-z0-9]*|mxfp\d|nvfp4|bf16|fp16|f16|fp8|dwq|optiq|gguf|mlx|4bit|8bit)+$", "", name) name = re.sub(r"-(gguf|mlx)$", "", name) return name class Harvester: def __init__(self, settings: Settings, db, registry, jobs: JobRunner, downloader): self.settings = settings self.db = db self.registry = registry self.jobs = jobs self.downloader = downloader self._preferred_runtime = "mlx" def _api(self): from huggingface_hub import HfApi return HfApi(token=self.settings.hf_token or os.environ.get("HF_TOKEN") or None) async def start_scan(self, options: dict, actor: str | None = None) -> Job: payload = { "runtimes": options.get("runtimes") or ["mlx", "gguf"], "authors": options.get("authors") or None, "limit_per_author": int(options.get("limit_per_author") or 150), "min_downloads": int(options.get("min_downloads") or 500), "max_ram_gb": float(options.get("max_ram_gb") or 0) or None, "families": options.get("families") or None, "tasks": options.get("tasks") or None, "search": options.get("search") or None, } await self.db.audit("harvest.scan", actor=actor, detail=payload) async def run(job: Job): return await self._scan(job, payload) return self.jobs.submit("harvest", "Harvest Hugging Face", payload, run) async def _scan(self, job: Job, opt: dict) -> dict: api = self._api() budget, absolute = await self.registry.budgets() self._preferred_runtime = str(await self.db.get_setting("preferred_runtime", "mlx") or "mlx") max_ram = opt["max_ram_gb"] or budget installed = {r["repository"] for r in await self.db.fetchall("SELECT repository FROM models WHERE installed=1 AND repository IS NOT NULL")} installed_keys = set() for r in await self.db.fetchall("SELECT name, tags, repository FROM models WHERE installed=1"): installed_keys.add(_base_key(r["name"], json.loads(r["tags"] or "[]"))) if r["repository"]: installed_keys.add(_base_key(r["repository"], [])) sources: list[tuple[str, str]] = [] for rt in opt["runtimes"]: authors = opt["authors"] or DEFAULT_SOURCES.get(rt, []) for a in authors: sources.append((rt, a)) candidates: list[dict] = [] seen: set[str] = set() total_sources = max(1, len(sources)) for si, (rt, author) in enumerate(sources): if job.cancelled: break self.jobs.update(job, progress=0.05 + 0.6 * si / total_sources, stage=f"listing {author} ({rt})") kwargs: dict[str, Any] = {"author": author, "sort": "downloads", "limit": opt["limit_per_author"], "expand": ["downloads", "likes", "tags", "pipeline_tag", "lastModified", "config", "safetensors", "gated"]} if rt == "gguf": kwargs["filter"] = "gguf" if opt["search"]: kwargs["search"] = opt["search"] try: infos = await asyncio.to_thread(lambda: list(api.list_models(**kwargs))) except Exception as e: log.warning("list_models %s failed: %s", author, e) self.jobs.update(job, warning=f"{author}: {str(e)[:120]}") continue for info in infos: rid = info.id if rid in seen: continue seen.add(rid) name = rid.split("/")[-1] if EXCLUDE.search(name): continue tags = list(info.tags or []) pipeline = getattr(info, "pipeline_tag", None) if pipeline in ("automatic-speech-recognition", "text-to-speech", "text-to-image", "text-to-video", "audio-to-audio"): continue downloads = getattr(info, "downloads", 0) or 0 if downloads < opt["min_downloads"]: continue if getattr(info, "gated", False): continue candidates.append({"repo": rid, "runtime": "llamacpp" if rt == "gguf" else "mlx", "tags": tags, "pipeline": pipeline, "downloads": downloads, "likes": getattr(info, "likes", 0) or 0, "last_modified": str(getattr(info, "last_modified", "") or ""), "config": getattr(info, "config", None) or {}, "safetensors": getattr(info, "safetensors", None)}) # family/task filters fams = set(f.lower() for f in (opt["families"] or [])) tasks = set(t.lower() for t in (opt["tasks"] or [])) rows: list[dict] = [] n = len(candidates) for i, c in enumerate(candidates): if job.cancelled: break if i % 10 == 0: self.jobs.update(job, progress=0.65 + 0.3 * i / max(1, n), stage=f"evaluating {i}/{n}") name = c["repo"].split("/")[-1] family = formats.guess_family(name, (c["config"] or {}).get("model_type")) if fams and family not in fams: continue task_flags = {k: bool(p.search(name)) for k, p in TASK_HINTS.items()} if c["pipeline"] in ("feature-extraction", "sentence-similarity"): task_flags["embedding"] = True if c["pipeline"] == "text-ranking": task_flags["reranker"] = True if c["pipeline"] == "image-text-to-text": task_flags["vision"] = True task = ("reranker" if task_flags["reranker"] else "embedding" if task_flags["embedding"] else "vision" if task_flags["vision"] else "coding" if task_flags["coding"] else "reasoning" if task_flags["reasoning"] else "general") if tasks and task not in tasks and not (task_flags.get(next(iter(tasks), ""), False)): continue try: row = await self._evaluate(c, name, family, task, task_flags, budget, absolute, max_ram) except Exception as e: # one bad repo must not abort the harvest log.warning("harvest: skipping %s (%s: %s)", c["repo"], type(e).__name__, e) continue if not row: continue row["installed"] = int(c["repo"] in installed or row["base_model"] in installed_keys) rows.append(row) # dedupe: same base model + runtime -> keep the best quantization for the budget rows.sort(key=lambda r: (-r["score"])) best_by_key: dict[str, dict] = {} for r in rows: key = f"{r['runtime']}::{r['base_model']}" if key in best_by_key: r["duplicate_of"] = best_by_key[key]["repo_id"] else: best_by_key[key] = r now = time.time() await self.db.execute("DELETE FROM harvest_candidates WHERE selected=0 AND dismissed=0") for r in rows: await self.db.execute( "INSERT INTO harvest_candidates(repo_id, runtime, family, base_model, name, task, quantization, parameter_count, " "download_bytes, estimated_ram_gb, size_class, compatibility_status, compatibility_reason, downloads, likes, last_modified, " "files, duplicate_of, installed, score, scanned_at, raw) VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?) " "ON CONFLICT(repo_id) DO UPDATE SET runtime=excluded.runtime, family=excluded.family, base_model=excluded.base_model, " "name=excluded.name, task=excluded.task, quantization=excluded.quantization, parameter_count=excluded.parameter_count, " "download_bytes=excluded.download_bytes, estimated_ram_gb=excluded.estimated_ram_gb, size_class=excluded.size_class, " "compatibility_status=excluded.compatibility_status, compatibility_reason=excluded.compatibility_reason, downloads=excluded.downloads, " "likes=excluded.likes, last_modified=excluded.last_modified, files=excluded.files, duplicate_of=excluded.duplicate_of, " "installed=excluded.installed, score=excluded.score, scanned_at=excluded.scanned_at, raw=excluded.raw", (r["repo_id"], r["runtime"], r["family"], r["base_model"], r["name"], r["task"], r["quantization"], r["parameter_count"], r["download_bytes"], r["estimated_ram_gb"], r["size_class"], r["compatibility_status"], r["compatibility_reason"], r["downloads"], r["likes"], r["last_modified"], json.dumps(r["files"]), r.get("duplicate_of"), r["installed"], r["score"], now, json.dumps({"tags": r["tags"][:30], "task_flags": r["task_flags"]}))) summary = {"sources": len(sources), "listed": len(candidates), "candidates": len(rows), "unique": len(best_by_key), "compatible": sum(1 for r in rows if r["compatibility_status"] in (compat.COMPATIBLE, compat.RESTRICTED)), "at": now} return summary async def _evaluate(self, c: dict, name: str, family: str, task: str, flags: dict, budget: float, absolute: float, max_ram: float) -> dict | None: cfg = formats.parse_hf_config(c["config"]) if c["config"] else {} quant, bits = None, None download_bytes = 0 weights_bytes = 0 files: list[dict] = [] if c["runtime"] == "mlx": quant, bits = cfg.get("quantization"), cfg.get("quant_bits") if not quant: quant, bits = formats.parse_quant_from_name(name) st = c.get("safetensors") total_params = None if st is not None: total_params = getattr(st, "total", None) or (st.get("total") if isinstance(st, dict) else None) params_by_dtype = getattr(st, "parameters", None) or (st.get("parameters") if isinstance(st, dict) else None) if params_by_dtype and isinstance(params_by_dtype, dict): # Hugging Face reports quantized uint32 tensors by *unpacked* parameter count # (an 8B 4-bit repo shows U32 ≈ 8.0e9): bytes = params × bits / 8, plus ~12 % for scales/biases. for dt, cnt in params_by_dtype.items(): if dt in ("U32", "I32") and bits: weights_bytes += int(cnt * bits / 8 * 1.12) elif dt in ("F16", "BF16"): weights_bytes += int(cnt * 2) elif dt in ("F32",): weights_bytes += int(cnt * 4) elif dt in ("U8", "I8", "F8_E4M3", "F8_E5M2"): weights_bytes += int(cnt) else: weights_bytes += int(cnt * 2) pc, ac = formats.parse_param_count_from_name(name) if bits and pc and not weights_bytes: weights_bytes = int(pc * bits / 8 * 1.06) if not bits and not quant: quant, bits = "bf16", 16 if pc and not weights_bytes: weights_bytes = pc * 2 if bits and total_params and bits < 16 and weights_bytes: # 'total' counts packed uint32 elements as 1 param; approximate real param count pass param_count = pc or (int(total_params) if total_params else None) or (int(weights_bytes * 8 / bits) if bits and weights_bytes else None) download_bytes = int(weights_bytes * 1.02) if not weights_bytes: return None else: # GGUF: need the file list; use the lightweight repo tree (sizes) — one API call per repo try: api = self._api() tree = await asyncio.to_thread(lambda: list(api.list_repo_tree(c["repo"], recursive=False))) except Exception: return None files_all = [{"path": getattr(f, "path", ""), "size": getattr(f, "size", 0) or 0} for f in tree] sel = pick_gguf_files(files_all) if not sel: return None files = sel weights_bytes = sum(f["size"] for f in sel if "mmproj" not in f["path"].lower()) download_bytes = sum(f["size"] for f in sel) quant, bits = formats.parse_quant_from_name(sel[0]["path"].rsplit(".", 1)[0]) pc, ac = formats.parse_param_count_from_name(name) param_count = pc or (int(weights_bytes * 8 / bits) if bits else None) kv = kv_bytes_per_token(cfg.get("n_layers"), cfg.get("n_kv_heads"), cfg.get("head_dim"), 16, cfg.get("full_attention_layers"), cfg.get("sliding_window")) if not kv and param_count: # No config in the listing: ~130 KB/token for an 8B dense model, scaling gently; MoE models # (e.g. 30B-A3B) have KV sized like their active parameters, not their total. _, active = formats.parse_param_count_from_name(name) ref = active or param_count kv = int(130_000 * max(0.3, (ref / 8e9)) ** 0.5) vision = flags.get("vision", False) or bool(cfg.get("vision")) embedding = flags.get("embedding", False) reranker = flags.get("reranker", False) from .models.scanner import llamacpp_available comp = compat.evaluate(runtime=c["runtime"], weights_bytes=weights_bytes, kv_per_token=kv, max_context=cfg.get("max_context") or 32768, model_type=cfg.get("model_type") or ("llama" if c["runtime"] == "llamacpp" else None), architecture=(cfg.get("architectures") or [None])[0], vision=vision, embedding=embedding, reranker=reranker, budget_gb=min(budget, max_ram), absolute_gb=absolute, llamacpp_available=llamacpp_available(self.settings.llama_server_bin), quant_bits=bits, weights_file=files[0]["path"] if files else None) if c["runtime"] == "mlx" and comp.status == compat.INCOMPATIBLE and "Architecture" in comp.reason and not cfg.get("model_type"): comp = compat.Compatibility(compat.EXPERIMENTAL, "Architecture unknown (no config in listing); verify before download.", True, comp.estimated_ram_gb, comp.recommended_context, comp.estimate) if comp.status == compat.INCOMPATIBLE or comp.status == compat.NOT_RECOMMENDED: return None # quantization quality policy pol = self._quant_policy(param_count, bits) score = self._score(c, comp, bits, param_count, pol, task, self._preferred_runtime) return { "repo_id": c["repo"], "runtime": c["runtime"], "family": family, "base_model": _base_key(c["repo"], c["tags"]), "name": name, "task": task, "quantization": quant, "parameter_count": param_count, "download_bytes": download_bytes, "estimated_ram_gb": comp.estimated_ram_gb, "size_class": formats.size_class(comp.estimated_ram_gb, budget), "compatibility_status": comp.status, "compatibility_reason": comp.reason, "downloads": c["downloads"], "likes": c["likes"], "last_modified": c["last_modified"], "files": files, "score": score, "tags": c["tags"], "task_flags": flags, "quant_policy": pol, } @staticmethod def _quant_policy(params: int | None, bits: float | None) -> str: """ok | low | too_low relative to the size-based preference table.""" if not params or not bits: return "unknown" b = params / 1e9 if b <= 8: return "ok" if bits >= 6 else "low" if bits >= 4 else "too_low" if b <= 20: return "ok" if bits >= 5 else "low" if bits >= 4 else "too_low" if b <= 40: return "ok" if bits >= 4 else "low" if bits >= 3 else "too_low" if b <= 80: return "ok" if bits >= 4 else "low" if bits >= 3 else "too_low" return "ok" if bits >= 3 else "too_low" @staticmethod def _score(c: dict, comp, bits, params, pol: str, task: str, preferred_runtime: str = "mlx") -> float: import math s = math.log10(max(10, c["downloads"])) * 10 if c["runtime"] == preferred_runtime: s += 14 # Apple Silicon: the preferred runtime wins ties against more-downloaded GGUF mirrors s += math.log10(max(1, c["likes"])) * 3 s += {"ok": 15, "low": 5, "too_low": -20, "unknown": 0}[pol] s += {compat.COMPATIBLE: 10, compat.RESTRICTED: 4, compat.EXPERIMENTAL: -5}.get(comp.status, 0) if params: b = params / 1e9 s += min(15, b / 2) # bigger is (usually) better, capped if c["last_modified"] and c["last_modified"][:4].isdigit(): year = int(c["last_modified"][:4]) s += (year - 2024) * 4 return round(s, 2) # ------------------------------------------------------------------ queue async def candidates(self, *, include_duplicates: bool = False, task: str | None = None, runtime: str | None = None, family: str | None = None, size_class: str | None = None, q: str | None = None, limit: int = 300) -> list[dict]: sql = "SELECT * FROM harvest_candidates WHERE dismissed=0" params: list[Any] = [] if not include_duplicates: sql += " AND duplicate_of IS NULL" if task: sql += " AND task=?" params.append(task) if runtime: sql += " AND runtime=?" params.append(runtime) if family: sql += " AND family=?" params.append(family) if size_class: sql += " AND size_class=?" params.append(size_class) if q: sql += " AND (repo_id LIKE ? OR base_model LIKE ?)" params += [f"%{q}%", f"%{q}%"] sql += " ORDER BY selected DESC, score DESC LIMIT ?" params.append(limit) rows = await self.db.fetchall(sql, params) for r in rows: for k in ("files", "raw"): if r.get(k): try: r[k] = json.loads(r[k]) except Exception: pass return rows async def select(self, repo_id: str, selected: bool) -> None: await self.db.execute("UPDATE harvest_candidates SET selected=? WHERE repo_id=?", (int(selected), repo_id)) async def dismiss(self, repo_id: str) -> None: await self.db.execute("UPDATE harvest_candidates SET dismissed=1, selected=0 WHERE repo_id=?", (repo_id,)) async def queue_selected(self, actor: str | None = None) -> list[dict]: rows = await self.db.fetchall("SELECT * FROM harvest_candidates WHERE selected=1 AND installed=0 ORDER BY score DESC") out = [] for r in rows: try: job = await self.downloader.start_download(r["repo_id"], r["quantization"] if r["runtime"] == "llamacpp" else None, actor=actor) out.append({"repo": r["repo_id"], "job": job.id}) await self.db.execute("UPDATE harvest_candidates SET selected=0 WHERE repo_id=?", (r["repo_id"],)) except Exception as e: out.append({"repo": r["repo_id"], "error": str(e)}) return out async def suggest_starter(self) -> list[dict]: """A curated slot list (small/medium/large general, coding, reasoning, vision, embedding, reranker) filled from the latest harvest, best score first.""" rows = await self.candidates(limit=1000) def p(r): return (r["parameter_count"] or 0) / 1e9 def g(r): return r["estimated_ram_gb"] or 0 slots = { "small general": lambda r: r["task"] in ("general", "vision") and 3 <= p(r) <= 12 and g(r) < 12, "small coding": lambda r: r["task"] == "coding" and 3 <= p(r) <= 12, "small reasoning": lambda r: r["task"] == "reasoning" and 3 <= p(r) <= 12, "medium general": lambda r: r["task"] in ("general", "vision") and 12 < p(r) <= 32 and g(r) < 26, "medium coding": lambda r: r["task"] == "coding" and 12 < p(r) <= 40, "large general": lambda r: r["task"] in ("general", "vision") and p(r) > 24 and 20 <= g(r) <= 45, "large reasoning": lambda r: r["task"] == "reasoning" and p(r) > 12 and g(r) <= 45, "vision": lambda r: r["task"] == "vision" and p(r) >= 7, "embedding": lambda r: r["task"] == "embedding", "reranker": lambda r: r["task"] == "reranker", } out = [] used = set() for slot, pred in slots.items(): pick = next((r for r in rows if pred(r) and r["repo_id"] not in used and r["compatibility_status"] in (compat.COMPATIBLE, compat.RESTRICTED)), None) if pick: used.add(pick["repo_id"]) out.append({"slot": slot, "candidate": pick}) return out