1"""Model Harvester: explore Hugging Face for MLX/GGUF models that genuinely fit this machine,2deduplicate them and propose a download queue. Never downloads blindly."""34from __future__ import annotations56import asyncio7import json8import logging9import os10import re11import time12from typing import Any1314from .config import Settings15from .downloads import TRUSTED_GGUF_AUTHORS, pick_gguf_files16from .jobs import Job, JobRunner17from .models import compat, formats18from .models.estimator import kv_bytes_per_token1920log = logging.getLogger("llm_api.harvester")21GB = 1024**32223DEFAULT_SOURCES = {24 "mlx": ["mlx-community"],25 "gguf": ["unsloth", "bartowski", "ggml-org", "lmstudio-community"],26}2728TASK_HINTS = {29 "coding": re.compile(r"coder|code|devstral|codestral|starcoder|deepseek-coder|kat-coder", re.I),30 "reasoning": re.compile(r"thinking|reason|r1|qwq|deepseek-r|magistral|-think", re.I),31 "vision": re.compile(r"-vl|vision|pixtral|llava|gemma-?3|gemma-?4|qwen3\.[5-8]|paligemma|-omni|ocr", re.I),32 "embedding": re.compile(r"embed|e5-|bge|gte-|minilm|nomic", re.I),33 "reranker": re.compile(r"rerank", re.I),34 "multilingual": re.compile(r"multilingual|qwen|gemma|aya|mistral", re.I),35}3637EXCLUDE = re.compile(r"uncensored|abliterated|heretic|nsfw|erotic|roleplay|-rp-|waifu|distill(ed)?-mlx-4bit-claude|"38 r"whisper|parakeet|tts|kokoro|orpheus|-asr|speech|audio|text-to-image|qwen-image|flux|stable-diffusion|"39 r"sd-|wan2|video|roformer|-base(-|$)|pretrain|draft|eagle|mtp-", re.I)404142def _base_key(repo: str, tags: list[str]) -> str:43 bm = next((t.split(":", 2)[2] for t in tags if t.startswith("base_model:quantized:")), None)44 if not bm:45 bm = next((t.split(":", 1)[1] for t in tags if t.startswith("base_model:") and "finetune:" not in t), None)46 if bm:47 return bm.lower()48 name = repo.split("/")[-1].lower()49 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)50 name = re.sub(r"-(gguf|mlx)$", "", name)51 return name525354class Harvester:55 def __init__(self, settings: Settings, db, registry, jobs: JobRunner, downloader):56 self.settings = settings57 self.db = db58 self.registry = registry59 self.jobs = jobs60 self.downloader = downloader61 self._preferred_runtime = "mlx"6263 def _api(self):64 from huggingface_hub import HfApi65 return HfApi(token=self.settings.hf_token or os.environ.get("HF_TOKEN") or None)6667 async def start_scan(self, options: dict, actor: str | None = None) -> Job:68 payload = {69 "runtimes": options.get("runtimes") or ["mlx", "gguf"],70 "authors": options.get("authors") or None,71 "limit_per_author": int(options.get("limit_per_author") or 150),72 "min_downloads": int(options.get("min_downloads") or 500),73 "max_ram_gb": float(options.get("max_ram_gb") or 0) or None,74 "families": options.get("families") or None,75 "tasks": options.get("tasks") or None,76 "search": options.get("search") or None,77 }78 await self.db.audit("harvest.scan", actor=actor, detail=payload)7980 async def run(job: Job):81 return await self._scan(job, payload)8283 return self.jobs.submit("harvest", "Harvest Hugging Face", payload, run)8485 async def _scan(self, job: Job, opt: dict) -> dict:86 api = self._api()87 budget, absolute = await self.registry.budgets()88 self._preferred_runtime = str(await self.db.get_setting("preferred_runtime", "mlx") or "mlx")89 max_ram = opt["max_ram_gb"] or budget90 installed = {r["repository"] for r in await self.db.fetchall("SELECT repository FROM models WHERE installed=1 AND repository IS NOT NULL")}91 installed_keys = set()92 for r in await self.db.fetchall("SELECT name, tags FROM models WHERE installed=1"):93 installed_keys.add(_base_key(r["name"], json.loads(r["tags"] or "[]")))94 sources: list[tuple[str, str]] = []95 for rt in opt["runtimes"]:96 authors = opt["authors"] or DEFAULT_SOURCES.get(rt, [])97 for a in authors:98 sources.append((rt, a))99 candidates: list[dict] = []100 seen: set[str] = set()101 total_sources = max(1, len(sources))102 for si, (rt, author) in enumerate(sources):103 if job.cancelled:104 break105 self.jobs.update(job, progress=0.05 + 0.6 * si / total_sources, stage=f"listing {author} ({rt})")106 kwargs: dict[str, Any] = {"author": author, "sort": "downloads", "limit": opt["limit_per_author"],107 "expand": ["downloads", "likes", "tags", "pipeline_tag", "lastModified", "config", "safetensors", "gated"]}108 if rt == "gguf":109 kwargs["filter"] = "gguf"110 if opt["search"]:111 kwargs["search"] = opt["search"]112 try:113 infos = await asyncio.to_thread(lambda: list(api.list_models(**kwargs)))114 except Exception as e:115 log.warning("list_models %s failed: %s", author, e)116 self.jobs.update(job, warning=f"{author}: {str(e)[:120]}")117 continue118 for info in infos:119 rid = info.id120 if rid in seen:121 continue122 seen.add(rid)123 name = rid.split("/")[-1]124 if EXCLUDE.search(name):125 continue126 tags = list(info.tags or [])127 pipeline = getattr(info, "pipeline_tag", None)128 if pipeline in ("automatic-speech-recognition", "text-to-speech", "text-to-image", "text-to-video", "audio-to-audio"):129 continue130 downloads = getattr(info, "downloads", 0) or 0131 if downloads < opt["min_downloads"]:132 continue133 if getattr(info, "gated", False):134 continue135 candidates.append({"repo": rid, "runtime": "llamacpp" if rt == "gguf" else "mlx", "tags": tags, "pipeline": pipeline,136 "downloads": downloads, "likes": getattr(info, "likes", 0) or 0,137 "last_modified": str(getattr(info, "last_modified", "") or ""),138 "config": getattr(info, "config", None) or {},139 "safetensors": getattr(info, "safetensors", None)})140 # family/task filters141 fams = set(f.lower() for f in (opt["families"] or []))142 tasks = set(t.lower() for t in (opt["tasks"] or []))143 rows: list[dict] = []144 n = len(candidates)145 for i, c in enumerate(candidates):146 if job.cancelled:147 break148 if i % 10 == 0:149 self.jobs.update(job, progress=0.65 + 0.3 * i / max(1, n), stage=f"evaluating {i}/{n}")150 name = c["repo"].split("/")[-1]151 family = formats.guess_family(name, (c["config"] or {}).get("model_type"))152 if fams and family not in fams:153 continue154 task_flags = {k: bool(p.search(name)) for k, p in TASK_HINTS.items()}155 if c["pipeline"] in ("feature-extraction", "sentence-similarity"):156 task_flags["embedding"] = True157 if c["pipeline"] == "text-ranking":158 task_flags["reranker"] = True159 if c["pipeline"] == "image-text-to-text":160 task_flags["vision"] = True161 task = ("reranker" if task_flags["reranker"] else "embedding" if task_flags["embedding"] else162 "vision" if task_flags["vision"] else "coding" if task_flags["coding"] else163 "reasoning" if task_flags["reasoning"] else "general")164 if tasks and task not in tasks and not (task_flags.get(next(iter(tasks), ""), False)):165 continue166 try:167 row = await self._evaluate(c, name, family, task, task_flags, budget, absolute, max_ram)168 except Exception as e: # one bad repo must not abort the harvest169 log.warning("harvest: skipping %s (%s: %s)", c["repo"], type(e).__name__, e)170 continue171 if not row:172 continue173 row["installed"] = int(c["repo"] in installed or row["base_model"] in installed_keys)174 rows.append(row)175 # dedupe: same base model + runtime -> keep the best quantization for the budget176 rows.sort(key=lambda r: (-r["score"]))177 best_by_key: dict[str, dict] = {}178 for r in rows:179 key = f"{r['runtime']}::{r['base_model']}"180 if key in best_by_key:181 r["duplicate_of"] = best_by_key[key]["repo_id"]182 else:183 best_by_key[key] = r184 now = time.time()185 await self.db.execute("DELETE FROM harvest_candidates WHERE selected=0 AND dismissed=0")186 for r in rows:187 await self.db.execute(188 "INSERT INTO harvest_candidates(repo_id, runtime, family, base_model, name, task, quantization, parameter_count, "189 "download_bytes, estimated_ram_gb, size_class, compatibility_status, compatibility_reason, downloads, likes, last_modified, "190 "files, duplicate_of, installed, score, scanned_at, raw) VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?) "191 "ON CONFLICT(repo_id) DO UPDATE SET runtime=excluded.runtime, family=excluded.family, base_model=excluded.base_model, "192 "name=excluded.name, task=excluded.task, quantization=excluded.quantization, parameter_count=excluded.parameter_count, "193 "download_bytes=excluded.download_bytes, estimated_ram_gb=excluded.estimated_ram_gb, size_class=excluded.size_class, "194 "compatibility_status=excluded.compatibility_status, compatibility_reason=excluded.compatibility_reason, downloads=excluded.downloads, "195 "likes=excluded.likes, last_modified=excluded.last_modified, files=excluded.files, duplicate_of=excluded.duplicate_of, "196 "installed=excluded.installed, score=excluded.score, scanned_at=excluded.scanned_at, raw=excluded.raw",197 (r["repo_id"], r["runtime"], r["family"], r["base_model"], r["name"], r["task"], r["quantization"], r["parameter_count"],198 r["download_bytes"], r["estimated_ram_gb"], r["size_class"], r["compatibility_status"], r["compatibility_reason"],199 r["downloads"], r["likes"], r["last_modified"], json.dumps(r["files"]), r.get("duplicate_of"), r["installed"],200 r["score"], now, json.dumps({"tags": r["tags"][:30], "task_flags": r["task_flags"]})))201 summary = {"sources": len(sources), "listed": len(candidates), "candidates": len(rows),202 "unique": len(best_by_key), "compatible": sum(1 for r in rows if r["compatibility_status"] in (compat.COMPATIBLE, compat.RESTRICTED)),203 "at": now}204 return summary205206 async def _evaluate(self, c: dict, name: str, family: str, task: str, flags: dict, budget: float, absolute: float,207 max_ram: float) -> dict | None:208 cfg = formats.parse_hf_config(c["config"]) if c["config"] else {}209 quant, bits = None, None210 download_bytes = 0211 weights_bytes = 0212 files: list[dict] = []213 if c["runtime"] == "mlx":214 quant, bits = cfg.get("quantization"), cfg.get("quant_bits")215 if not quant:216 quant, bits = formats.parse_quant_from_name(name)217 st = c.get("safetensors")218 total_params = None219 if st is not None:220 total_params = getattr(st, "total", None) or (st.get("total") if isinstance(st, dict) else None)221 params_by_dtype = getattr(st, "parameters", None) or (st.get("parameters") if isinstance(st, dict) else None)222 if params_by_dtype and isinstance(params_by_dtype, dict):223 # packed uint32 quantized weights: count * 32 / bits224 for dt, cnt in params_by_dtype.items():225 if dt in ("U32", "I32") and bits:226 weights_bytes += int(cnt * 4)227 elif dt in ("F16", "BF16"):228 weights_bytes += int(cnt * 2)229 elif dt in ("F32",):230 weights_bytes += int(cnt * 4)231 elif dt in ("U8", "I8", "F8_E4M3", "F8_E5M2"):232 weights_bytes += int(cnt)233 else:234 weights_bytes += int(cnt * 2)235 pc, ac = formats.parse_param_count_from_name(name)236 if bits and pc and not weights_bytes:237 weights_bytes = int(pc * bits / 8 * 1.06)238 if not bits and not quant:239 quant, bits = "bf16", 16240 if pc and not weights_bytes:241 weights_bytes = pc * 2242 if bits and total_params and bits < 16 and weights_bytes:243 # 'total' counts packed uint32 elements as 1 param; approximate real param count244 pass245 param_count = pc or (int(weights_bytes * 8 / bits) if bits and weights_bytes else None)246 download_bytes = int(weights_bytes * 1.02)247 if not weights_bytes:248 return None249 else:250 # GGUF: need the file list; use the lightweight repo tree (sizes) — one API call per repo251 try:252 api = self._api()253 tree = await asyncio.to_thread(lambda: list(api.list_repo_tree(c["repo"], recursive=False)))254 except Exception:255 return None256 files_all = [{"path": getattr(f, "path", ""), "size": getattr(f, "size", 0) or 0} for f in tree]257 sel = pick_gguf_files(files_all)258 if not sel:259 return None260 files = sel261 weights_bytes = sum(f["size"] for f in sel if "mmproj" not in f["path"].lower())262 download_bytes = sum(f["size"] for f in sel)263 quant, bits = formats.parse_quant_from_name(sel[0]["path"].rsplit(".", 1)[0])264 pc, ac = formats.parse_param_count_from_name(name)265 param_count = pc or (int(weights_bytes * 8 / bits) if bits else None)266 kv = kv_bytes_per_token(cfg.get("n_layers"), cfg.get("n_kv_heads"), cfg.get("head_dim"), 16,267 cfg.get("full_attention_layers"), cfg.get("sliding_window"))268 if not kv and param_count:269 # No config in the listing: ~130 KB/token for an 8B dense model, scaling gently; MoE models270 # (e.g. 30B-A3B) have KV sized like their active parameters, not their total.271 _, active = formats.parse_param_count_from_name(name)272 ref = active or param_count273 kv = int(130_000 * max(0.3, (ref / 8e9)) ** 0.5)274 vision = flags.get("vision", False) or bool(cfg.get("vision"))275 embedding = flags.get("embedding", False)276 reranker = flags.get("reranker", False)277 from .models.scanner import llamacpp_available278 comp = compat.evaluate(runtime=c["runtime"], weights_bytes=weights_bytes, kv_per_token=kv, max_context=cfg.get("max_context") or 32768,279 model_type=cfg.get("model_type") or ("llama" if c["runtime"] == "llamacpp" else None),280 architecture=(cfg.get("architectures") or [None])[0], vision=vision, embedding=embedding,281 reranker=reranker, budget_gb=min(budget, max_ram), absolute_gb=absolute,282 llamacpp_available=llamacpp_available(self.settings.llama_server_bin), quant_bits=bits,283 weights_file=files[0]["path"] if files else None)284 if c["runtime"] == "mlx" and comp.status == compat.INCOMPATIBLE and "Architecture" in comp.reason and not cfg.get("model_type"):285 comp = compat.Compatibility(compat.EXPERIMENTAL, "Architecture unknown (no config in listing); verify before download.",286 True, comp.estimated_ram_gb, comp.recommended_context, comp.estimate)287 if comp.status == compat.INCOMPATIBLE or comp.status == compat.NOT_RECOMMENDED:288 return None289 # quantization quality policy290 pol = self._quant_policy(param_count, bits)291 score = self._score(c, comp, bits, param_count, pol, task, self._preferred_runtime)292 return {293 "repo_id": c["repo"], "runtime": c["runtime"], "family": family, "base_model": _base_key(c["repo"], c["tags"]),294 "name": name, "task": task, "quantization": quant, "parameter_count": param_count, "download_bytes": download_bytes,295 "estimated_ram_gb": comp.estimated_ram_gb, "size_class": formats.size_class(comp.estimated_ram_gb, budget),296 "compatibility_status": comp.status, "compatibility_reason": comp.reason, "downloads": c["downloads"],297 "likes": c["likes"], "last_modified": c["last_modified"], "files": files, "score": score, "tags": c["tags"],298 "task_flags": flags, "quant_policy": pol,299 }300301 @staticmethod302 def _quant_policy(params: int | None, bits: float | None) -> str:303 """ok | low | too_low relative to the size-based preference table."""304 if not params or not bits:305 return "unknown"306 b = params / 1e9307 if b <= 8:308 return "ok" if bits >= 6 else "low" if bits >= 4 else "too_low"309 if b <= 20:310 return "ok" if bits >= 5 else "low" if bits >= 4 else "too_low"311 if b <= 40:312 return "ok" if bits >= 4 else "low" if bits >= 3 else "too_low"313 if b <= 80:314 return "ok" if bits >= 4 else "low" if bits >= 3 else "too_low"315 return "ok" if bits >= 3 else "too_low"316317 @staticmethod318 def _score(c: dict, comp, bits, params, pol: str, task: str, preferred_runtime: str = "mlx") -> float:319 import math320 s = math.log10(max(10, c["downloads"])) * 10321 if c["runtime"] == preferred_runtime:322 s += 14 # Apple Silicon: the preferred runtime wins ties against more-downloaded GGUF mirrors323 s += math.log10(max(1, c["likes"])) * 3324 s += {"ok": 15, "low": 5, "too_low": -20, "unknown": 0}[pol]325 s += {compat.COMPATIBLE: 10, compat.RESTRICTED: 4, compat.EXPERIMENTAL: -5}.get(comp.status, 0)326 if params:327 b = params / 1e9328 s += min(15, b / 2) # bigger is (usually) better, capped329 if c["last_modified"] and c["last_modified"][:4].isdigit():330 year = int(c["last_modified"][:4])331 s += (year - 2024) * 4332 return round(s, 2)333334 # ------------------------------------------------------------------ queue335 async def candidates(self, *, include_duplicates: bool = False, task: str | None = None, runtime: str | None = None,336 family: str | None = None, size_class: str | None = None, q: str | None = None, limit: int = 300) -> list[dict]:337 sql = "SELECT * FROM harvest_candidates WHERE dismissed=0"338 params: list[Any] = []339 if not include_duplicates:340 sql += " AND duplicate_of IS NULL"341 if task:342 sql += " AND task=?"343 params.append(task)344 if runtime:345 sql += " AND runtime=?"346 params.append(runtime)347 if family:348 sql += " AND family=?"349 params.append(family)350 if size_class:351 sql += " AND size_class=?"352 params.append(size_class)353 if q:354 sql += " AND (repo_id LIKE ? OR base_model LIKE ?)"355 params += [f"%{q}%", f"%{q}%"]356 sql += " ORDER BY selected DESC, score DESC LIMIT ?"357 params.append(limit)358 rows = await self.db.fetchall(sql, params)359 for r in rows:360 for k in ("files", "raw"):361 if r.get(k):362 try:363 r[k] = json.loads(r[k])364 except Exception:365 pass366 return rows367368 async def select(self, repo_id: str, selected: bool) -> None:369 await self.db.execute("UPDATE harvest_candidates SET selected=? WHERE repo_id=?", (int(selected), repo_id))370371 async def dismiss(self, repo_id: str) -> None:372 await self.db.execute("UPDATE harvest_candidates SET dismissed=1, selected=0 WHERE repo_id=?", (repo_id,))373374 async def queue_selected(self, actor: str | None = None) -> list[dict]:375 rows = await self.db.fetchall("SELECT * FROM harvest_candidates WHERE selected=1 AND installed=0 ORDER BY score DESC")376 out = []377 for r in rows:378 try:379 job = await self.downloader.start_download(r["repo_id"], r["quantization"] if r["runtime"] == "llamacpp" else None, actor=actor)380 out.append({"repo": r["repo_id"], "job": job.id})381 await self.db.execute("UPDATE harvest_candidates SET selected=0 WHERE repo_id=?", (r["repo_id"],))382 except Exception as e:383 out.append({"repo": r["repo_id"], "error": str(e)})384 return out385386 async def suggest_starter(self) -> list[dict]:387 """A curated slot list (small/medium/large general, coding, reasoning, vision, embedding, reranker)388 filled from the latest harvest, best score first."""389 rows = await self.candidates(limit=1000)390 def p(r):391 return (r["parameter_count"] or 0) / 1e9392393 def g(r):394 return r["estimated_ram_gb"] or 0395396 slots = {397 "small general": lambda r: r["task"] in ("general", "vision") and 3 <= p(r) <= 12 and g(r) < 12,398 "small coding": lambda r: r["task"] == "coding" and 3 <= p(r) <= 12,399 "small reasoning": lambda r: r["task"] == "reasoning" and 3 <= p(r) <= 12,400 "medium general": lambda r: r["task"] in ("general", "vision") and 12 < p(r) <= 32 and g(r) < 26,401 "medium coding": lambda r: r["task"] == "coding" and 12 < p(r) <= 40,402 "large general": lambda r: r["task"] in ("general", "vision") and p(r) > 24 and 20 <= g(r) <= 45,403 "large reasoning": lambda r: r["task"] == "reasoning" and p(r) > 12 and g(r) <= 45,404 "vision": lambda r: r["task"] == "vision" and p(r) >= 7,405 "embedding": lambda r: r["task"] == "embedding",406 "reranker": lambda r: r["task"] == "reranker",407 }408 out = []409 used = set()410 for slot, pred in slots.items():411 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)412 if pick:413 used.add(pick["repo_id"])414 out.append({"slot": slot, "candidate": pick})415 return out416