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 = downloader6162 def _api(self):63 from huggingface_hub import HfApi64 return HfApi(token=self.settings.hf_token or os.environ.get("HF_TOKEN") or None)6566 async def start_scan(self, options: dict, actor: str | None = None) -> Job:67 payload = {68 "runtimes": options.get("runtimes") or ["mlx", "gguf"],69 "authors": options.get("authors") or None,70 "limit_per_author": int(options.get("limit_per_author") or 150),71 "min_downloads": int(options.get("min_downloads") or 500),72 "max_ram_gb": float(options.get("max_ram_gb") or 0) or None,73 "families": options.get("families") or None,74 "tasks": options.get("tasks") or None,75 "search": options.get("search") or None,76 }77 await self.db.audit("harvest.scan", actor=actor, detail=payload)7879 async def run(job: Job):80 return await self._scan(job, payload)8182 return self.jobs.submit("harvest", "Harvest Hugging Face", payload, run)8384 async def _scan(self, job: Job, opt: dict) -> dict:85 api = self._api()86 budget, absolute = await self.registry.budgets()87 max_ram = opt["max_ram_gb"] or budget88 installed = {r["repository"] for r in await self.db.fetchall("SELECT repository FROM models WHERE installed=1 AND repository IS NOT NULL")}89 installed_keys = set()90 for r in await self.db.fetchall("SELECT name, tags FROM models WHERE installed=1"):91 installed_keys.add(_base_key(r["name"], json.loads(r["tags"] or "[]")))92 sources: list[tuple[str, str]] = []93 for rt in opt["runtimes"]:94 authors = opt["authors"] or DEFAULT_SOURCES.get(rt, [])95 for a in authors:96 sources.append((rt, a))97 candidates: list[dict] = []98 seen: set[str] = set()99 total_sources = max(1, len(sources))100 for si, (rt, author) in enumerate(sources):101 if job.cancelled:102 break103 self.jobs.update(job, progress=0.05 + 0.6 * si / total_sources, stage=f"listing {author} ({rt})")104 kwargs: dict[str, Any] = {"author": author, "sort": "downloads", "limit": opt["limit_per_author"],105 "expand": ["downloads", "likes", "tags", "pipeline_tag", "lastModified", "config", "safetensors", "gated"]}106 if rt == "gguf":107 kwargs["filter"] = "gguf"108 if opt["search"]:109 kwargs["search"] = opt["search"]110 try:111 infos = await asyncio.to_thread(lambda: list(api.list_models(**kwargs)))112 except Exception as e:113 log.warning("list_models %s failed: %s", author, e)114 self.jobs.update(job, warning=f"{author}: {str(e)[:120]}")115 continue116 for info in infos:117 rid = info.id118 if rid in seen:119 continue120 seen.add(rid)121 name = rid.split("/")[-1]122 if EXCLUDE.search(name):123 continue124 tags = list(info.tags or [])125 pipeline = getattr(info, "pipeline_tag", None)126 if pipeline in ("automatic-speech-recognition", "text-to-speech", "text-to-image", "text-to-video", "audio-to-audio"):127 continue128 downloads = getattr(info, "downloads", 0) or 0129 if downloads < opt["min_downloads"]:130 continue131 if getattr(info, "gated", False):132 continue133 candidates.append({"repo": rid, "runtime": "llamacpp" if rt == "gguf" else "mlx", "tags": tags, "pipeline": pipeline,134 "downloads": downloads, "likes": getattr(info, "likes", 0) or 0,135 "last_modified": str(getattr(info, "last_modified", "") or ""),136 "config": getattr(info, "config", None) or {},137 "safetensors": getattr(info, "safetensors", None)})138 # family/task filters139 fams = set(f.lower() for f in (opt["families"] or []))140 tasks = set(t.lower() for t in (opt["tasks"] or []))141 rows: list[dict] = []142 n = len(candidates)143 for i, c in enumerate(candidates):144 if job.cancelled:145 break146 if i % 10 == 0:147 self.jobs.update(job, progress=0.65 + 0.3 * i / max(1, n), stage=f"evaluating {i}/{n}")148 name = c["repo"].split("/")[-1]149 family = formats.guess_family(name, (c["config"] or {}).get("model_type"))150 if fams and family not in fams:151 continue152 task_flags = {k: bool(p.search(name)) for k, p in TASK_HINTS.items()}153 if c["pipeline"] in ("feature-extraction", "sentence-similarity"):154 task_flags["embedding"] = True155 if c["pipeline"] == "text-ranking":156 task_flags["reranker"] = True157 if c["pipeline"] == "image-text-to-text":158 task_flags["vision"] = True159 task = ("reranker" if task_flags["reranker"] else "embedding" if task_flags["embedding"] else160 "vision" if task_flags["vision"] else "coding" if task_flags["coding"] else161 "reasoning" if task_flags["reasoning"] else "general")162 if tasks and task not in tasks and not (task_flags.get(next(iter(tasks), ""), False)):163 continue164 try:165 row = await self._evaluate(c, name, family, task, task_flags, budget, absolute, max_ram)166 except Exception as e: # one bad repo must not abort the harvest167 log.warning("harvest: skipping %s (%s: %s)", c["repo"], type(e).__name__, e)168 continue169 if not row:170 continue171 row["installed"] = int(c["repo"] in installed or row["base_model"] in installed_keys)172 rows.append(row)173 # dedupe: same base model + runtime -> keep the best quantization for the budget174 rows.sort(key=lambda r: (-r["score"]))175 best_by_key: dict[str, dict] = {}176 for r in rows:177 key = f"{r['runtime']}::{r['base_model']}"178 if key in best_by_key:179 r["duplicate_of"] = best_by_key[key]["repo_id"]180 else:181 best_by_key[key] = r182 now = time.time()183 await self.db.execute("DELETE FROM harvest_candidates WHERE selected=0 AND dismissed=0")184 for r in rows:185 await self.db.execute(186 "INSERT INTO harvest_candidates(repo_id, runtime, family, base_model, name, task, quantization, parameter_count, "187 "download_bytes, estimated_ram_gb, size_class, compatibility_status, compatibility_reason, downloads, likes, last_modified, "188 "files, duplicate_of, installed, score, scanned_at, raw) VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?) "189 "ON CONFLICT(repo_id) DO UPDATE SET runtime=excluded.runtime, family=excluded.family, base_model=excluded.base_model, "190 "name=excluded.name, task=excluded.task, quantization=excluded.quantization, parameter_count=excluded.parameter_count, "191 "download_bytes=excluded.download_bytes, estimated_ram_gb=excluded.estimated_ram_gb, size_class=excluded.size_class, "192 "compatibility_status=excluded.compatibility_status, compatibility_reason=excluded.compatibility_reason, downloads=excluded.downloads, "193 "likes=excluded.likes, last_modified=excluded.last_modified, files=excluded.files, duplicate_of=excluded.duplicate_of, "194 "installed=excluded.installed, score=excluded.score, scanned_at=excluded.scanned_at, raw=excluded.raw",195 (r["repo_id"], r["runtime"], r["family"], r["base_model"], r["name"], r["task"], r["quantization"], r["parameter_count"],196 r["download_bytes"], r["estimated_ram_gb"], r["size_class"], r["compatibility_status"], r["compatibility_reason"],197 r["downloads"], r["likes"], r["last_modified"], json.dumps(r["files"]), r.get("duplicate_of"), r["installed"],198 r["score"], now, json.dumps({"tags": r["tags"][:30], "task_flags": r["task_flags"]})))199 summary = {"sources": len(sources), "listed": len(candidates), "candidates": len(rows),200 "unique": len(best_by_key), "compatible": sum(1 for r in rows if r["compatibility_status"] in (compat.COMPATIBLE, compat.RESTRICTED)),201 "at": now}202 return summary203204 async def _evaluate(self, c: dict, name: str, family: str, task: str, flags: dict, budget: float, absolute: float,205 max_ram: float) -> dict | None:206 cfg = formats.parse_hf_config(c["config"]) if c["config"] else {}207 quant, bits = None, None208 download_bytes = 0209 weights_bytes = 0210 files: list[dict] = []211 if c["runtime"] == "mlx":212 quant, bits = cfg.get("quantization"), cfg.get("quant_bits")213 if not quant:214 quant, bits = formats.parse_quant_from_name(name)215 st = c.get("safetensors")216 total_params = None217 if st is not None:218 total_params = getattr(st, "total", None) or (st.get("total") if isinstance(st, dict) else None)219 params_by_dtype = getattr(st, "parameters", None) or (st.get("parameters") if isinstance(st, dict) else None)220 if params_by_dtype and isinstance(params_by_dtype, dict):221 # packed uint32 quantized weights: count * 32 / bits222 for dt, cnt in params_by_dtype.items():223 if dt in ("U32", "I32") and bits:224 weights_bytes += int(cnt * 4)225 elif dt in ("F16", "BF16"):226 weights_bytes += int(cnt * 2)227 elif dt in ("F32",):228 weights_bytes += int(cnt * 4)229 elif dt in ("U8", "I8", "F8_E4M3", "F8_E5M2"):230 weights_bytes += int(cnt)231 else:232 weights_bytes += int(cnt * 2)233 pc, ac = formats.parse_param_count_from_name(name)234 if bits and pc and not weights_bytes:235 weights_bytes = int(pc * bits / 8 * 1.06)236 if not bits and not quant:237 quant, bits = "bf16", 16238 if pc and not weights_bytes:239 weights_bytes = pc * 2240 if bits and total_params and bits < 16 and weights_bytes:241 # 'total' counts packed uint32 elements as 1 param; approximate real param count242 pass243 param_count = pc or (int(weights_bytes * 8 / bits) if bits and weights_bytes else None)244 download_bytes = int(weights_bytes * 1.02)245 if not weights_bytes:246 return None247 else:248 # GGUF: need the file list; use the lightweight repo tree (sizes) — one API call per repo249 try:250 api = self._api()251 tree = await asyncio.to_thread(lambda: list(api.list_repo_tree(c["repo"], recursive=False)))252 except Exception:253 return None254 files_all = [{"path": getattr(f, "path", ""), "size": getattr(f, "size", 0) or 0} for f in tree]255 sel = pick_gguf_files(files_all)256 if not sel:257 return None258 files = sel259 weights_bytes = sum(f["size"] for f in sel if "mmproj" not in f["path"].lower())260 download_bytes = sum(f["size"] for f in sel)261 quant, bits = formats.parse_quant_from_name(sel[0]["path"].rsplit(".", 1)[0])262 pc, ac = formats.parse_param_count_from_name(name)263 param_count = pc or (int(weights_bytes * 8 / bits) if bits else None)264 kv = kv_bytes_per_token(cfg.get("n_layers"), cfg.get("n_kv_heads"), cfg.get("head_dim"), 16,265 cfg.get("full_attention_layers"), cfg.get("sliding_window"))266 if not kv and param_count:267 kv = int(130_000 * (param_count / 8e9) ** 0.6)268 vision = flags.get("vision", False) or bool(cfg.get("vision"))269 embedding = flags.get("embedding", False)270 reranker = flags.get("reranker", False)271 from .models.scanner import llamacpp_available272 comp = compat.evaluate(runtime=c["runtime"], weights_bytes=weights_bytes, kv_per_token=kv, max_context=cfg.get("max_context") or 32768,273 model_type=cfg.get("model_type") or ("llama" if c["runtime"] == "llamacpp" else None),274 architecture=(cfg.get("architectures") or [None])[0], vision=vision, embedding=embedding,275 reranker=reranker, budget_gb=min(budget, max_ram), absolute_gb=absolute,276 llamacpp_available=llamacpp_available(self.settings.llama_server_bin), quant_bits=bits,277 weights_file=files[0]["path"] if files else None)278 if c["runtime"] == "mlx" and comp.status == compat.INCOMPATIBLE and "Architecture" in comp.reason and not cfg.get("model_type"):279 comp = compat.Compatibility(compat.EXPERIMENTAL, "Architecture unknown (no config in listing); verify before download.",280 True, comp.estimated_ram_gb, comp.recommended_context, comp.estimate)281 if comp.status == compat.INCOMPATIBLE or comp.status == compat.NOT_RECOMMENDED:282 return None283 # quantization quality policy284 pol = self._quant_policy(param_count, bits)285 score = self._score(c, comp, bits, param_count, pol, task)286 return {287 "repo_id": c["repo"], "runtime": c["runtime"], "family": family, "base_model": _base_key(c["repo"], c["tags"]),288 "name": name, "task": task, "quantization": quant, "parameter_count": param_count, "download_bytes": download_bytes,289 "estimated_ram_gb": comp.estimated_ram_gb, "size_class": formats.size_class(comp.estimated_ram_gb, budget),290 "compatibility_status": comp.status, "compatibility_reason": comp.reason, "downloads": c["downloads"],291 "likes": c["likes"], "last_modified": c["last_modified"], "files": files, "score": score, "tags": c["tags"],292 "task_flags": flags, "quant_policy": pol,293 }294295 @staticmethod296 def _quant_policy(params: int | None, bits: float | None) -> str:297 """ok | low | too_low relative to the size-based preference table."""298 if not params or not bits:299 return "unknown"300 b = params / 1e9301 if b <= 8:302 return "ok" if bits >= 6 else "low" if bits >= 4 else "too_low"303 if b <= 20:304 return "ok" if bits >= 5 else "low" if bits >= 4 else "too_low"305 if b <= 40:306 return "ok" if bits >= 4 else "low" if bits >= 3 else "too_low"307 if b <= 80:308 return "ok" if bits >= 4 else "low" if bits >= 3 else "too_low"309 return "ok" if bits >= 3 else "too_low"310311 @staticmethod312 def _score(c: dict, comp, bits, params, pol: str, task: str) -> float:313 import math314 s = math.log10(max(10, c["downloads"])) * 10315 s += math.log10(max(1, c["likes"])) * 3316 s += {"ok": 15, "low": 5, "too_low": -20, "unknown": 0}[pol]317 s += {compat.COMPATIBLE: 10, compat.RESTRICTED: 4, compat.EXPERIMENTAL: -5}.get(comp.status, 0)318 if params:319 b = params / 1e9320 s += min(15, b / 2) # bigger is (usually) better, capped321 if c["last_modified"] and c["last_modified"][:4].isdigit():322 year = int(c["last_modified"][:4])323 s += (year - 2024) * 4324 return round(s, 2)325326 # ------------------------------------------------------------------ queue327 async def candidates(self, *, include_duplicates: bool = False, task: str | None = None, runtime: str | None = None,328 family: str | None = None, size_class: str | None = None, q: str | None = None, limit: int = 300) -> list[dict]:329 sql = "SELECT * FROM harvest_candidates WHERE dismissed=0"330 params: list[Any] = []331 if not include_duplicates:332 sql += " AND duplicate_of IS NULL"333 if task:334 sql += " AND task=?"335 params.append(task)336 if runtime:337 sql += " AND runtime=?"338 params.append(runtime)339 if family:340 sql += " AND family=?"341 params.append(family)342 if size_class:343 sql += " AND size_class=?"344 params.append(size_class)345 if q:346 sql += " AND (repo_id LIKE ? OR base_model LIKE ?)"347 params += [f"%{q}%", f"%{q}%"]348 sql += " ORDER BY selected DESC, score DESC LIMIT ?"349 params.append(limit)350 rows = await self.db.fetchall(sql, params)351 for r in rows:352 for k in ("files", "raw"):353 if r.get(k):354 try:355 r[k] = json.loads(r[k])356 except Exception:357 pass358 return rows359360 async def select(self, repo_id: str, selected: bool) -> None:361 await self.db.execute("UPDATE harvest_candidates SET selected=? WHERE repo_id=?", (int(selected), repo_id))362363 async def dismiss(self, repo_id: str) -> None:364 await self.db.execute("UPDATE harvest_candidates SET dismissed=1, selected=0 WHERE repo_id=?", (repo_id,))365366 async def queue_selected(self, actor: str | None = None) -> list[dict]:367 rows = await self.db.fetchall("SELECT * FROM harvest_candidates WHERE selected=1 AND installed=0 ORDER BY score DESC")368 out = []369 for r in rows:370 try:371 job = await self.downloader.start_download(r["repo_id"], r["quantization"] if r["runtime"] == "llamacpp" else None, actor=actor)372 out.append({"repo": r["repo_id"], "job": job.id})373 await self.db.execute("UPDATE harvest_candidates SET selected=0 WHERE repo_id=?", (r["repo_id"],))374 except Exception as e:375 out.append({"repo": r["repo_id"], "error": str(e)})376 return out377378 async def suggest_starter(self) -> list[dict]:379 """A curated slot list (small/medium/large general, coding, reasoning, vision, embedding, reranker)380 filled from the latest harvest, best score first."""381 rows = await self.candidates(limit=1000)382 slots = {383 "small general": lambda r: r["task"] == "general" and (r["estimated_ram_gb"] or 0) < 8,384 "small coding": lambda r: r["task"] == "coding" and (r["estimated_ram_gb"] or 0) < 12,385 "small reasoning": lambda r: r["task"] == "reasoning" and (r["estimated_ram_gb"] or 0) < 12,386 "medium general": lambda r: r["task"] == "general" and 8 <= (r["estimated_ram_gb"] or 0) < 22,387 "medium coding": lambda r: r["task"] == "coding" and 12 <= (r["estimated_ram_gb"] or 0) < 25,388 "large general": lambda r: r["task"] == "general" and 22 <= (r["estimated_ram_gb"] or 0) <= 45,389 "large reasoning": lambda r: r["task"] == "reasoning" and 12 <= (r["estimated_ram_gb"] or 0) <= 45,390 "vision": lambda r: r["task"] == "vision",391 "embedding": lambda r: r["task"] == "embedding",392 "reranker": lambda r: r["task"] == "reranker",393 }394 out = []395 used = set()396 for slot, pred in slots.items():397 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)398 if pick:399 used.add(pick["repo_id"])400 out.append({"slot": slot, "candidate": pick})401 return out402