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.split("/")[-1].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, repository FROM models WHERE installed=1"):93 installed_keys.add(_base_key(r["name"], json.loads(r["tags"] or "[]")))94 if r["repository"]:95 installed_keys.add(_base_key(r["repository"], []))96 sources: list[tuple[str, str]] = []97 for rt in opt["runtimes"]:98 authors = opt["authors"] or DEFAULT_SOURCES.get(rt, [])99 for a in authors:100 sources.append((rt, a))101 candidates: list[dict] = []102 seen: set[str] = set()103 total_sources = max(1, len(sources))104 for si, (rt, author) in enumerate(sources):105 if job.cancelled:106 break107 self.jobs.update(job, progress=0.05 + 0.6 * si / total_sources, stage=f"listing {author} ({rt})")108 kwargs: dict[str, Any] = {"author": author, "sort": "downloads", "limit": opt["limit_per_author"],109 "expand": ["downloads", "likes", "tags", "pipeline_tag", "lastModified", "config", "safetensors", "gated"]}110 if rt == "gguf":111 kwargs["filter"] = "gguf"112 if opt["search"]:113 kwargs["search"] = opt["search"]114 try:115 infos = await asyncio.to_thread(lambda: list(api.list_models(**kwargs)))116 except Exception as e:117 log.warning("list_models %s failed: %s", author, e)118 self.jobs.update(job, warning=f"{author}: {str(e)[:120]}")119 continue120 for info in infos:121 rid = info.id122 if rid in seen:123 continue124 seen.add(rid)125 name = rid.split("/")[-1]126 if EXCLUDE.search(name):127 continue128 tags = list(info.tags or [])129 pipeline = getattr(info, "pipeline_tag", None)130 if pipeline in ("automatic-speech-recognition", "text-to-speech", "text-to-image", "text-to-video", "audio-to-audio"):131 continue132 downloads = getattr(info, "downloads", 0) or 0133 if downloads < opt["min_downloads"]:134 continue135 if getattr(info, "gated", False):136 continue137 candidates.append({"repo": rid, "runtime": "llamacpp" if rt == "gguf" else "mlx", "tags": tags, "pipeline": pipeline,138 "downloads": downloads, "likes": getattr(info, "likes", 0) or 0,139 "last_modified": str(getattr(info, "last_modified", "") or ""),140 "config": getattr(info, "config", None) or {},141 "safetensors": getattr(info, "safetensors", None)})142 # family/task filters143 fams = set(f.lower() for f in (opt["families"] or []))144 tasks = set(t.lower() for t in (opt["tasks"] or []))145 rows: list[dict] = []146 n = len(candidates)147 for i, c in enumerate(candidates):148 if job.cancelled:149 break150 if i % 10 == 0:151 self.jobs.update(job, progress=0.65 + 0.3 * i / max(1, n), stage=f"evaluating {i}/{n}")152 name = c["repo"].split("/")[-1]153 family = formats.guess_family(name, (c["config"] or {}).get("model_type"))154 if fams and family not in fams:155 continue156 task_flags = {k: bool(p.search(name)) for k, p in TASK_HINTS.items()}157 if c["pipeline"] in ("feature-extraction", "sentence-similarity"):158 task_flags["embedding"] = True159 if c["pipeline"] == "text-ranking":160 task_flags["reranker"] = True161 if c["pipeline"] == "image-text-to-text":162 task_flags["vision"] = True163 task = ("reranker" if task_flags["reranker"] else "embedding" if task_flags["embedding"] else164 "vision" if task_flags["vision"] else "coding" if task_flags["coding"] else165 "reasoning" if task_flags["reasoning"] else "general")166 if tasks and task not in tasks and not (task_flags.get(next(iter(tasks), ""), False)):167 continue168 try:169 row = await self._evaluate(c, name, family, task, task_flags, budget, absolute, max_ram)170 except Exception as e: # one bad repo must not abort the harvest171 log.warning("harvest: skipping %s (%s: %s)", c["repo"], type(e).__name__, e)172 continue173 if not row:174 continue175 row["installed"] = int(c["repo"] in installed or row["base_model"] in installed_keys)176 rows.append(row)177 # dedupe: same base model + runtime -> keep the best quantization for the budget178 rows.sort(key=lambda r: (-r["score"]))179 best_by_key: dict[str, dict] = {}180 for r in rows:181 key = f"{r['runtime']}::{r['base_model']}"182 if key in best_by_key:183 r["duplicate_of"] = best_by_key[key]["repo_id"]184 else:185 best_by_key[key] = r186 now = time.time()187 await self.db.execute("DELETE FROM harvest_candidates WHERE selected=0 AND dismissed=0")188 for r in rows:189 await self.db.execute(190 "INSERT INTO harvest_candidates(repo_id, runtime, family, base_model, name, task, quantization, parameter_count, "191 "download_bytes, estimated_ram_gb, size_class, compatibility_status, compatibility_reason, downloads, likes, last_modified, "192 "files, duplicate_of, installed, score, scanned_at, raw) VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?) "193 "ON CONFLICT(repo_id) DO UPDATE SET runtime=excluded.runtime, family=excluded.family, base_model=excluded.base_model, "194 "name=excluded.name, task=excluded.task, quantization=excluded.quantization, parameter_count=excluded.parameter_count, "195 "download_bytes=excluded.download_bytes, estimated_ram_gb=excluded.estimated_ram_gb, size_class=excluded.size_class, "196 "compatibility_status=excluded.compatibility_status, compatibility_reason=excluded.compatibility_reason, downloads=excluded.downloads, "197 "likes=excluded.likes, last_modified=excluded.last_modified, files=excluded.files, duplicate_of=excluded.duplicate_of, "198 "installed=excluded.installed, score=excluded.score, scanned_at=excluded.scanned_at, raw=excluded.raw",199 (r["repo_id"], r["runtime"], r["family"], r["base_model"], r["name"], r["task"], r["quantization"], r["parameter_count"],200 r["download_bytes"], r["estimated_ram_gb"], r["size_class"], r["compatibility_status"], r["compatibility_reason"],201 r["downloads"], r["likes"], r["last_modified"], json.dumps(r["files"]), r.get("duplicate_of"), r["installed"],202 r["score"], now, json.dumps({"tags": r["tags"][:30], "task_flags": r["task_flags"]})))203 summary = {"sources": len(sources), "listed": len(candidates), "candidates": len(rows),204 "unique": len(best_by_key), "compatible": sum(1 for r in rows if r["compatibility_status"] in (compat.COMPATIBLE, compat.RESTRICTED)),205 "at": now}206 return summary207208 async def _evaluate(self, c: dict, name: str, family: str, task: str, flags: dict, budget: float, absolute: float,209 max_ram: float) -> dict | None:210 cfg = formats.parse_hf_config(c["config"]) if c["config"] else {}211 quant, bits = None, None212 download_bytes = 0213 weights_bytes = 0214 files: list[dict] = []215 if c["runtime"] == "mlx":216 quant, bits = cfg.get("quantization"), cfg.get("quant_bits")217 if not quant:218 quant, bits = formats.parse_quant_from_name(name)219 st = c.get("safetensors")220 total_params = None221 if st is not None:222 total_params = getattr(st, "total", None) or (st.get("total") if isinstance(st, dict) else None)223 params_by_dtype = getattr(st, "parameters", None) or (st.get("parameters") if isinstance(st, dict) else None)224 if params_by_dtype and isinstance(params_by_dtype, dict):225 # Hugging Face reports quantized uint32 tensors by *unpacked* parameter count226 # (an 8B 4-bit repo shows U32 ≈ 8.0e9): bytes = params × bits / 8, plus ~12 % for scales/biases.227 for dt, cnt in params_by_dtype.items():228 if dt in ("U32", "I32") and bits:229 weights_bytes += int(cnt * bits / 8 * 1.12)230 elif dt in ("F16", "BF16"):231 weights_bytes += int(cnt * 2)232 elif dt in ("F32",):233 weights_bytes += int(cnt * 4)234 elif dt in ("U8", "I8", "F8_E4M3", "F8_E5M2"):235 weights_bytes += int(cnt)236 else:237 weights_bytes += int(cnt * 2)238 pc, ac = formats.parse_param_count_from_name(name)239 if bits and pc and not weights_bytes:240 weights_bytes = int(pc * bits / 8 * 1.06)241 if not bits and not quant:242 quant, bits = "bf16", 16243 if pc and not weights_bytes:244 weights_bytes = pc * 2245 if bits and total_params and bits < 16 and weights_bytes:246 # 'total' counts packed uint32 elements as 1 param; approximate real param count247 pass248 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)249 download_bytes = int(weights_bytes * 1.02)250 if not weights_bytes:251 return None252 else:253 # GGUF: need the file list; use the lightweight repo tree (sizes) — one API call per repo254 try:255 api = self._api()256 tree = await asyncio.to_thread(lambda: list(api.list_repo_tree(c["repo"], recursive=False)))257 except Exception:258 return None259 files_all = [{"path": getattr(f, "path", ""), "size": getattr(f, "size", 0) or 0} for f in tree]260 sel = pick_gguf_files(files_all)261 if not sel:262 return None263 files = sel264 weights_bytes = sum(f["size"] for f in sel if "mmproj" not in f["path"].lower())265 download_bytes = sum(f["size"] for f in sel)266 quant, bits = formats.parse_quant_from_name(sel[0]["path"].rsplit(".", 1)[0])267 pc, ac = formats.parse_param_count_from_name(name)268 param_count = pc or (int(weights_bytes * 8 / bits) if bits else None)269 kv = kv_bytes_per_token(cfg.get("n_layers"), cfg.get("n_kv_heads"), cfg.get("head_dim"), 16,270 cfg.get("full_attention_layers"), cfg.get("sliding_window"))271 if not kv and param_count:272 # No config in the listing: ~130 KB/token for an 8B dense model, scaling gently; MoE models273 # (e.g. 30B-A3B) have KV sized like their active parameters, not their total.274 _, active = formats.parse_param_count_from_name(name)275 ref = active or param_count276 kv = int(130_000 * max(0.3, (ref / 8e9)) ** 0.5)277 vision = flags.get("vision", False) or bool(cfg.get("vision"))278 embedding = flags.get("embedding", False)279 reranker = flags.get("reranker", False)280 from .models.scanner import llamacpp_available281 comp = compat.evaluate(runtime=c["runtime"], weights_bytes=weights_bytes, kv_per_token=kv, max_context=cfg.get("max_context") or 32768,282 model_type=cfg.get("model_type") or ("llama" if c["runtime"] == "llamacpp" else None),283 architecture=(cfg.get("architectures") or [None])[0], vision=vision, embedding=embedding,284 reranker=reranker, budget_gb=min(budget, max_ram), absolute_gb=absolute,285 llamacpp_available=llamacpp_available(self.settings.llama_server_bin), quant_bits=bits,286 weights_file=files[0]["path"] if files else None)287 if c["runtime"] == "mlx" and comp.status == compat.INCOMPATIBLE and "Architecture" in comp.reason and not cfg.get("model_type"):288 comp = compat.Compatibility(compat.EXPERIMENTAL, "Architecture unknown (no config in listing); verify before download.",289 True, comp.estimated_ram_gb, comp.recommended_context, comp.estimate)290 if comp.status == compat.INCOMPATIBLE or comp.status == compat.NOT_RECOMMENDED:291 return None292 # quantization quality policy293 pol = self._quant_policy(param_count, bits)294 score = self._score(c, comp, bits, param_count, pol, task, self._preferred_runtime)295 return {296 "repo_id": c["repo"], "runtime": c["runtime"], "family": family, "base_model": _base_key(c["repo"], c["tags"]),297 "name": name, "task": task, "quantization": quant, "parameter_count": param_count, "download_bytes": download_bytes,298 "estimated_ram_gb": comp.estimated_ram_gb, "size_class": formats.size_class(comp.estimated_ram_gb, budget),299 "compatibility_status": comp.status, "compatibility_reason": comp.reason, "downloads": c["downloads"],300 "likes": c["likes"], "last_modified": c["last_modified"], "files": files, "score": score, "tags": c["tags"],301 "task_flags": flags, "quant_policy": pol,302 }303304 @staticmethod305 def _quant_policy(params: int | None, bits: float | None) -> str:306 """ok | low | too_low relative to the size-based preference table."""307 if not params or not bits:308 return "unknown"309 b = params / 1e9310 if b <= 8:311 return "ok" if bits >= 6 else "low" if bits >= 4 else "too_low"312 if b <= 20:313 return "ok" if bits >= 5 else "low" if bits >= 4 else "too_low"314 if b <= 40:315 return "ok" if bits >= 4 else "low" if bits >= 3 else "too_low"316 if b <= 80:317 return "ok" if bits >= 4 else "low" if bits >= 3 else "too_low"318 return "ok" if bits >= 3 else "too_low"319320 @staticmethod321 def _score(c: dict, comp, bits, params, pol: str, task: str, preferred_runtime: str = "mlx") -> float:322 import math323 s = math.log10(max(10, c["downloads"])) * 10324 if c["runtime"] == preferred_runtime:325 s += 14 # Apple Silicon: the preferred runtime wins ties against more-downloaded GGUF mirrors326 s += math.log10(max(1, c["likes"])) * 3327 s += {"ok": 15, "low": 5, "too_low": -20, "unknown": 0}[pol]328 s += {compat.COMPATIBLE: 10, compat.RESTRICTED: 4, compat.EXPERIMENTAL: -5}.get(comp.status, 0)329 if params:330 b = params / 1e9331 s += min(15, b / 2) # bigger is (usually) better, capped332 if c["last_modified"] and c["last_modified"][:4].isdigit():333 year = int(c["last_modified"][:4])334 s += (year - 2024) * 4335 return round(s, 2)336337 # ------------------------------------------------------------------ queue338 async def candidates(self, *, include_duplicates: bool = False, task: str | None = None, runtime: str | None = None,339 family: str | None = None, size_class: str | None = None, q: str | None = None, limit: int = 300) -> list[dict]:340 sql = "SELECT * FROM harvest_candidates WHERE dismissed=0"341 params: list[Any] = []342 if not include_duplicates:343 sql += " AND duplicate_of IS NULL"344 if task:345 sql += " AND task=?"346 params.append(task)347 if runtime:348 sql += " AND runtime=?"349 params.append(runtime)350 if family:351 sql += " AND family=?"352 params.append(family)353 if size_class:354 sql += " AND size_class=?"355 params.append(size_class)356 if q:357 sql += " AND (repo_id LIKE ? OR base_model LIKE ?)"358 params += [f"%{q}%", f"%{q}%"]359 sql += " ORDER BY selected DESC, score DESC LIMIT ?"360 params.append(limit)361 rows = await self.db.fetchall(sql, params)362 for r in rows:363 for k in ("files", "raw"):364 if r.get(k):365 try:366 r[k] = json.loads(r[k])367 except Exception:368 pass369 return rows370371 async def select(self, repo_id: str, selected: bool) -> None:372 await self.db.execute("UPDATE harvest_candidates SET selected=? WHERE repo_id=?", (int(selected), repo_id))373374 async def dismiss(self, repo_id: str) -> None:375 await self.db.execute("UPDATE harvest_candidates SET dismissed=1, selected=0 WHERE repo_id=?", (repo_id,))376377 async def queue_selected(self, actor: str | None = None) -> list[dict]:378 rows = await self.db.fetchall("SELECT * FROM harvest_candidates WHERE selected=1 AND installed=0 ORDER BY score DESC")379 out = []380 for r in rows:381 try:382 job = await self.downloader.start_download(r["repo_id"], r["quantization"] if r["runtime"] == "llamacpp" else None, actor=actor)383 out.append({"repo": r["repo_id"], "job": job.id})384 await self.db.execute("UPDATE harvest_candidates SET selected=0 WHERE repo_id=?", (r["repo_id"],))385 except Exception as e:386 out.append({"repo": r["repo_id"], "error": str(e)})387 return out388389 async def suggest_starter(self) -> list[dict]:390 """A curated slot list (small/medium/large general, coding, reasoning, vision, embedding, reranker)391 filled from the latest harvest, best score first."""392 rows = await self.candidates(limit=1000)393 def p(r):394 return (r["parameter_count"] or 0) / 1e9395396 def g(r):397 return r["estimated_ram_gb"] or 0398399 slots = {400 "small general": lambda r: r["task"] in ("general", "vision") and 3 <= p(r) <= 12 and g(r) < 12,401 "small coding": lambda r: r["task"] == "coding" and 3 <= p(r) <= 12,402 "small reasoning": lambda r: r["task"] == "reasoning" and 3 <= p(r) <= 12,403 "medium general": lambda r: r["task"] in ("general", "vision") and 12 < p(r) <= 32 and g(r) < 26,404 "medium coding": lambda r: r["task"] == "coding" and 12 < p(r) <= 40,405 "large general": lambda r: r["task"] in ("general", "vision") and p(r) > 24 and 20 <= g(r) <= 45,406 "large reasoning": lambda r: r["task"] == "reasoning" and p(r) > 12 and g(r) <= 45,407 "vision": lambda r: r["task"] == "vision" and p(r) >= 7,408 "embedding": lambda r: r["task"] == "embedding",409 "reranker": lambda r: r["task"] == "reranker",410 }411 out = []412 used = set()413 for slot, pred in slots.items():414 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)415 if pick:416 used.add(pick["repo_id"])417 out.append({"slot": slot, "candidate": pick})418 return out419