SPB Git forge

spb/ai-atlas

Public
41commits 1branches 0releases
4.6 MBsize
maindefault branch
12 days agolast push
HTML 77.2% TypeScript 10.5% Python 9.6% JavaScript 2.5%
14.4 KB · 256 lines python
Raw Blame History
1"""Find-a-model (API 1.1): deterministic matching of canonical models against use-case criteria. No LLM, no composite "winner" score —2each match lists the observed facts (`why`) that satisfied each criterion and the observed dimensions; sorting is by number of satisfied3criteria, then by the best benchmark rank the model holds anywhere, then by release date.45Rules (documented in `RULES`, returned to the client as `filters_applied`):6    coding        current result on a coding-category benchmark, or capabilities/modalities mention code7    reasoning     attributes.reasoning is true, or a current result on a reasoning/math benchmark8    agentic       current result on an agentic benchmark, or tool_calling is true9    long_context  context_length ≥ context_min (default 128 000)10    vision        modalities (any direction) include image11    low_cost      cheapest current output price ≤ max_output_price, or in the bottom quartile of all cheapest output prices12    local         weights downloadable (open-weights / open-source / restricted-weights) and the estimated footprint fits memory_gb13    embeddings    modalities include embedding (or pipeline tag says so)14    chat          text modality (or unspecified) and reachable through ≥ 1 current deployment or downloadable weights15"""16from __future__ import annotations1718from collections import defaultdict19from typing import Any2021from sqlalchemy.ext.asyncio import AsyncConnection2223from aiatlas.db import fetch_all24from aiatlas.ontology.licenses import LICENSES, normalize_license25from aiatlas.services import hardware_fit as hf26from aiatlas.services.frontier import all_primary_groups, rank_rows2728USE_CASES = ("coding", "reasoning", "agentic", "long_context", "vision", "low_cost", "local", "embeddings", "chat")29DOWNLOADABLE = ("open-weights", "open-source", "restricted-weights", "restricted")30CODING_CATEGORIES = ("coding",)31REASONING_CATEGORIES = ("reasoning", "math")32AGENTIC_CATEGORIES = ("agentic",)33RULES = {34    "coding": "current result on a coding-category benchmark, or capabilities/modalities mention code",35    "reasoning": "attributes.reasoning = true, or a current result on a reasoning/math benchmark",36    "agentic": "current result on an agentic benchmark, or tool_calling = true",37    "long_context": "context_length ≥ context_min (default 128 000 tokens)",38    "vision": "modalities include image",39    "low_cost": "cheapest current output price ≤ max_output_price, or in the bottom quartile of all models' cheapest output prices",40    "local": "weights downloadable and the ESTIMATED footprint (services.hardware_fit) fits memory_gb at the requested quantisation",41    "embeddings": "modalities include embedding",42    "chat": "text modality and reachable (≥ 1 current deployment or downloadable weights)",43    "deployment": "local → downloadable weights · api → ≥ 1 current deployment · any → no constraint",44    "license": "commercial → licence key known in the ontology with commercial_use = true",45    "openness": "attributes.openness in the requested categories",46    "max_input_price / max_output_price": "cheapest current provider price (USD per 1M tokens) ≤ the bound",47    "modalities": "every requested modality present in modalities / modalities_input / modalities_output",48}495051def _modalities(attrs: dict[str, Any]) -> set[str]:52    out: set[str] = set()53    for k in ("modalities", "modalities_input", "modalities_output"):54        v = attrs.get(k)55        if isinstance(v, list):56            out |= {str(x).strip().lower() for x in v}57        elif isinstance(v, str):58            out |= {x.strip().lower() for x in v.split(",") if x.strip()}59    if attrs.get("vision") is True:60        out.add("image")61    if attrs.get("audio") is True:62        out.add("audio")63    return {("image" if m in ("vision", "images") else "document" if m == "pdf" else m) for m in out}646566def _num(v: Any) -> float | None:67    if isinstance(v, bool) or v is None:68        return None69    try:70        return float(v)71    except (TypeError, ValueError):72        return None737475def _mentions_code(attrs: dict[str, Any]) -> bool:76    caps = attrs.get("capabilities")77    text = " ".join(str(c) for c in caps).lower() if isinstance(caps, list) else str(caps or "").lower()78    return "code" in text or "code" in _modalities(attrs) or str(attrs.get("pipeline_tag") or "").lower().startswith("text-generation") and "coder" in str(attrs.get("family") or "").lower()798081async def find_models(conn: AsyncConnection, *, use_case: str | None, deployment: str = "any", memory_gb: float | None = None, quant: str = "4bit",82                      context_min: int | None = None, license: str = "any", openness: list[str] | None = None, max_input_price: float | None = None,83                      max_output_price: float | None = None, modalities: list[str] | None = None, limit: int = 30) -> dict[str, Any]:84    models = await fetch_all(conn, """85        select e.id, e.slug, e.canonical_name, e.attributes, e.organization_id, o.slug as org_slug, o.canonical_name as org_name, e.first_seen_at86        from entities e left join entities o on o.id = e.organization_id where e.entity_type = 'model' and e.merged_into is null""")87    prices = await fetch_all(conn, """88        select p.model_id, min(p.input_per_mtok) filter (where p.input_per_mtok > 0) as min_input, min(p.output_per_mtok) filter (where p.output_per_mtok > 0) as min_output,89               count(distinct p.provider_id) as providers90        from prices p where p.valid_to is null group by 1""")91    price_by = {p["model_id"]: p for p in prices}92    groups = await all_primary_groups(conn)93    ranks: dict[str, dict[str, int]] = defaultdict(dict)       # model_id → {benchmark slug: rank}94    categories: dict[str, set[str]] = defaultdict(set)         # model_id → benchmark categories with a current result95    for g in groups.values():96        cat = (g["benchmark"].get("category") or "").lower()97        for r in rank_rows(g["rows"], g["higher_is_better"]):98            ranks[r["model_id"]][g["benchmark"]["slug"]] = r["rank"]99            if cat:100                categories[r["model_id"]].add(cat)101    outputs = sorted(float(p["min_output"]) for p in prices if p["min_output"] is not None)102    bottom_quartile = outputs[len(outputs) // 4] if outputs else None103104    ctx_min = context_min or 128_000105    quant = hf.normalize_quant(quant)106    want_mods = {m.strip().lower() for m in (modalities or []) if m.strip()}107    open_set = set(openness or [])108    matches: list[dict[str, Any]] = []109    for m in models:110        attrs = m["attributes"] or {}111        why: list[str] = []112        failed = False113        mods = _modalities(attrs)114        pr = price_by.get(m["id"])115        cheapest_out = _num(pr["min_output"]) if pr else None116        cheapest_in = _num(pr["min_input"]) if pr else None117        providers = int(pr["providers"]) if pr else 0118        opn = str(attrs.get("openness") or "")119        downloadable = opn in DOWNLOADABLE120        ctx = _num(attrs.get("context_length"))121        cats = categories.get(m["id"], set())122        estimated_fit: dict[str, Any] | None = None123124        # --- hard filters125        if deployment == "local" and not downloadable:126            continue127        if deployment == "api" and providers == 0:128            continue129        if open_set and opn not in open_set:130            continue131        if license == "commercial":132            key = attrs.get("license_key") or normalize_license(attrs.get("license"))133            info = LICENSES.get(key) if key else None134            if not info or info.commercial_use is not True:135                continue136            why.append(f"licence {info.label} allows commercial use")137        if max_input_price is not None:138            if cheapest_in is None or cheapest_in > max_input_price:139                continue140            why.append(f"cheapest input ${cheapest_in:g}/M ≤ ${max_input_price:g}")141        if max_output_price is not None and use_case != "low_cost":142            if cheapest_out is None or cheapest_out > max_output_price:143                continue144            why.append(f"cheapest output ${cheapest_out:g}/M ≤ ${max_output_price:g}")145        if want_mods and not want_mods <= mods:146            continue147        if want_mods:148            why.append("modalities include " + ", ".join(sorted(want_mods)))149        if memory_gb is not None and (use_case == "local" or deployment == "local"):150            estimated_fit = hf.fit_detailed(attrs, memory_gb, quant=quant, context=min(int(ctx or 8192), 8192))151            if estimated_fit is None or not estimated_fit["fits"]:152                continue153            why.append(f"estimated {estimated_fit['estimated_memory_gb']} GB at {quant} fits {memory_gb:g} GB (estimate)")154155        # --- use-case rule (one criterion; `why` records the observed fact)156        if use_case == "coding":157            if cats & set(CODING_CATEGORIES):158                why.append("current result on a coding benchmark: " + ", ".join(sorted(b for b in ranks[m["id"]] if groups_cat(groups, b) == "coding")[:4]))159            elif _mentions_code(attrs):160                why.append("capabilities/modalities mention code")161            else:162                failed = True163        elif use_case == "reasoning":164            if attrs.get("reasoning") is True:165                why.append("attributes.reasoning = true")166            elif cats & set(REASONING_CATEGORIES):167                why.append("current result on a reasoning/math benchmark")168            else:169                failed = True170        elif use_case == "agentic":171            if cats & set(AGENTIC_CATEGORIES):172                why.append("current result on an agentic benchmark")173            elif attrs.get("tool_calling") is True:174                why.append("tool_calling = true")175            else:176                failed = True177        elif use_case == "long_context":178            if ctx is not None and ctx >= ctx_min:179                why.append(f"context_length {int(ctx):,} ≥ {ctx_min:,}")180            else:181                failed = True182        elif use_case == "vision":183            if "image" in mods:184                why.append("modalities include image")185            else:186                failed = True187        elif use_case == "low_cost":188            bound = max_output_price if max_output_price is not None else bottom_quartile189            if cheapest_out is not None and bound is not None and cheapest_out <= bound:190                why.append(f"cheapest output ${cheapest_out:g}/M ≤ ${bound:g}/M " + ("(bound)" if max_output_price is not None else "(bottom quartile of observed prices)"))191            else:192                failed = True193        elif use_case == "local":194            if not downloadable:195                failed = True196            else:197                why.append(f"openness {opn}: weights downloadable")198                if memory_gb is None:199                    estimated_fit = hf.fit_detailed(attrs, 64.0, quant=quant, context=8192)200                    if estimated_fit is not None:201                        why.append(f"estimated {estimated_fit['estimated_memory_gb']} GB at {quant} (reference 64 GB device, estimate)")202        elif use_case == "embeddings":203            if "embedding" in mods or "embedding" in str(attrs.get("pipeline_tag") or "").lower():204                why.append("modalities include embedding")205            else:206                failed = True207        elif use_case == "chat":208            if (not mods or "text" in mods) and (providers > 0 or downloadable):209                why.append("text modality and reachable (" + (f"{providers} providers" if providers else "downloadable weights") + ")")210            else:211                failed = True212        if failed:213            continue214        if context_min is not None and use_case != "long_context":215            if ctx is None or ctx < context_min:216                continue217            why.append(f"context_length {int(ctx):,} ≥ {context_min:,}")218        best_rank = min(ranks[m["id"]].values()) if ranks.get(m["id"]) else None219        matches.append({220            "model": {"id": m["id"], "slug": m["slug"], "name": m["canonical_name"], "entity_type": "model",221                      "organization": {"id": m["organization_id"], "slug": m["org_slug"], "name": m["org_name"]} if m["organization_id"] else None},222            "why": why,223            "observed": {"context_length": ctx, "parameter_count": _num(attrs.get("parameter_count")), "openness": opn or None, "license": attrs.get("license"),224                         "license_key": attrs.get("license_key") or normalize_license(attrs.get("license")), "modalities": sorted(mods), "reasoning": attrs.get("reasoning"),225                         "tool_calling": attrs.get("tool_calling"), "release_date": attrs.get("release_date"), "cheapest_input_per_mtok": cheapest_in,226                         "cheapest_output_per_mtok": cheapest_out, "providers": providers, "benchmark_ranks": dict(sorted(ranks.get(m["id"], {}).items())), "best_rank": best_rank},227            **({"estimated_fit": estimated_fit} if estimated_fit else {}),228            "_sort": (-len(why), best_rank if best_rank is not None else 10_000, -(_date_ord(attrs.get("release_date")))),229        })230    matches.sort(key=lambda x: x["_sort"])231    for x in matches:232        x.pop("_sort", None)233    return {"matches": matches[:limit], "total": len(matches),234            "filters_applied": {k: v for k, v in {"use_case": use_case, "deployment": deployment, "memory_gb": memory_gb, "quant": quant if memory_gb is not None or use_case == "local" else None,235                                                  "context_min": context_min, "license": license, "openness": openness or None, "max_input_price": max_input_price,236                                                  "max_output_price": max_output_price, "modalities": sorted(want_mods) or None}.items() if v not in (None, "any")},237            "rules": {k: RULES[k] for k in RULES if k == use_case or k not in USE_CASES},238            "note": "Deterministic filters over observed attributes, current prices and current benchmark results of canonical models; sorted by number of satisfied "239                    "criteria, then best benchmark rank, then release date. No composite score. Hardware fit is an estimate."}240241242def groups_cat(groups: dict[str, dict[str, Any]], slug: str) -> str:243    for g in groups.values():244        if g["benchmark"]["slug"] == slug:245            return (g["benchmark"].get("category") or "").lower()246    return ""247248249def _date_ord(v: Any) -> int:250    s = str(v or "")251    digits = "".join(ch for ch in s[:10] if ch.isdigit())252    return int(digits.ljust(8, "0")[:8]) if digits else 0253254255__all__ = ["RULES", "USE_CASES", "find_models"]256