| """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 [])]
|
|
|
| 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
|
|
|
|
|
| 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:
|
| 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}"
|
| )
|
|
|
