blood-test-explainer / src /model_paths.py
Codex
Default extraction to the fine-tuned MiniCPM-V Hub checkpoint.
c1a5fb4
Raw
History Blame Contribute Delete
3.53 kB
"""Resolve MiniCPM-V Transformers weights: local disk first, Hub download fallback."""
from __future__ import annotations
import os
from dataclasses import dataclass
from pathlib import Path
# Fine-tuned MiniCPM-V 4.6 for lab extraction (MedReason SFT).
DEFAULT_HF_REPO = "build-small-hackathon/blood-test-minicpmv-4_6-medreason"
# Upstream OpenBMB base — eval baselines and training scripts.
BASE_HF_REPO = "openbmb/MiniCPM-V-4.6"
@dataclass(frozen=True)
class TransformersModelSource:
model_id: str
local_files_only: bool
origin: str
def project_root() -> Path:
return Path(__file__).resolve().parents[1]
def models_dir() -> Path:
raw = os.getenv("BTE_MODELS_DIR", "models").strip()
path = Path(raw)
if not path.is_absolute():
path = project_root() / path
return path.resolve()
def hub_cache_dir() -> Path:
cache = models_dir() / ".cache" / "huggingface" / "hub"
cache.mkdir(parents=True, exist_ok=True)
return cache
def apply_local_model_defaults() -> None:
"""Send Hugging Face downloads to the project models/ cache by default."""
models = models_dir()
os.environ.setdefault("BTE_MODELS_DIR", str(models))
hf_home = models / ".cache" / "huggingface"
hf_home.mkdir(parents=True, exist_ok=True)
os.environ.setdefault("HF_HOME", str(hf_home))
os.environ.setdefault("HUGGINGFACE_HUB_CACHE", str(hub_cache_dir()))
def resolve_transformers_model_source(model_id: str | None = None) -> TransformersModelSource:
"""Use a complete local checkpoint when present; otherwise download and load from Hub."""
configured = (model_id or os.getenv("ZEROGPU_MODEL_ID") or DEFAULT_HF_REPO).strip()
repo_id = DEFAULT_HF_REPO if configured.startswith(".") or configured.startswith("/") else configured
explicit = Path(configured).expanduser()
if explicit.is_dir() and is_transformers_model_dir(explicit):
return TransformersModelSource(str(explicit.resolve()), True, "local-dir")
for candidate in (
models_dir() / "MiniCPM-V-4.6",
models_dir() / "openbmb" / "MiniCPM-V-4.6",
models_dir() / repo_id.split("/", 1)[-1],
):
if is_transformers_model_dir(candidate):
return TransformersModelSource(str(candidate.resolve()), True, "local-dir")
for cache_root in (hub_cache_dir(), Path.home() / ".cache" / "huggingface" / "hub"):
snapshot = latest_complete_snapshot(repo_id, cache_root)
if snapshot:
label = "local-cache" if cache_root == hub_cache_dir() else "local-cache-global"
return TransformersModelSource(str(snapshot), True, label)
return TransformersModelSource(repo_id, False, "hub-download")
def is_transformers_model_dir(path: Path) -> bool:
if not path.is_dir() or not (path / "config.json").is_file():
return False
if list(path.glob("*.safetensors")) or list(path.glob("model*.bin")):
return True
return any(path.glob("model*.safetensors.index.json"))
def latest_complete_snapshot(repo_id: str, hub_cache: Path) -> Path | None:
if not hub_cache.is_dir():
return None
repo_dir = hub_cache / f"models--{repo_id.replace('/', '--')}" / "snapshots"
if not repo_dir.is_dir():
return None
snapshots = sorted(
(p for p in repo_dir.iterdir() if p.is_dir()),
key=lambda p: p.stat().st_mtime,
reverse=True,
)
for snapshot in snapshots:
if is_transformers_model_dir(snapshot):
return snapshot
return None