File size: 2,664 Bytes
a31a731 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | """Resolve a Kaggle-fitting HF base model (never GGUF)."""
from __future__ import annotations
import json
import os
from pathlib import Path
from typing import Any
DEFAULT_CANDIDATES = [
"Qwen/Qwen2.5-Coder-1.5B-Instruct",
"Qwen/Qwen2.5-Coder-0.5B-Instruct",
"Qwen/Qwen2.5-1.5B-Instruct",
"HuggingFaceTB/SmolLM2-1.7B-Instruct",
"google/gemma-2-2b-it",
]
def _probe(repo_id: str, token: str | None) -> bool:
from huggingface_hub import HfApi
api = HfApi(token=token)
try:
info = api.model_info(repo_id)
except Exception:
return False
siblings = [s.rfilename for s in (info.siblings or [])]
# Reject pure-GGUF repos
has_weight = any(
f.endswith((".safetensors", ".bin", ".pt"))
or f == "pytorch_model.bin.index.json"
or f.endswith(".safetensors.index.json")
for f in siblings
)
only_gguf = siblings and all(
f.endswith(".gguf") or f.endswith(".md") or "license" in f.lower() or f.startswith(".")
for f in siblings
)
if only_gguf and not has_weight:
return False
return has_weight or True # allow repos that list via remote code
def resolve_train_model(
preferred: list[str] | None = None,
*,
override: str | None = None,
token: str | None = None,
out_path: Path | None = None,
) -> dict[str, Any]:
token = token or os.getenv("HF_TOKEN")
override = override or os.getenv("TRAIN_MODEL_ID") or None
candidates = []
if override:
candidates.append(override)
candidates.extend(preferred or DEFAULT_CANDIDATES)
seen: set[str] = set()
errors: list[dict[str, str]] = []
for repo_id in candidates:
if repo_id in seen:
continue
seen.add(repo_id)
try:
if _probe(repo_id, token):
result = {
"train_model_id": repo_id,
"candidates_tried": list(seen),
"errors": errors,
}
if out_path:
out_path.parent.mkdir(parents=True, exist_ok=True)
out_path.write_text(json.dumps(result, indent=2), encoding="utf-8")
return result
except Exception as e: # noqa: BLE001
errors.append({"repo_id": repo_id, "error": f"{type(e).__name__}: {e}"})
continue
errors.append({"repo_id": repo_id, "error": "probe_failed"})
raise RuntimeError(
"No trainable HF model resolved. Tried: "
+ ", ".join(seen)
+ f" errors={errors}"
)
|