Fetch published config.json from Hub on load (enables download tracking)

halt_cot/transformers_backend.py

@@ -2,8 +2,9 @@
 
 from __future__ import annotations
 
-from dataclasses import replace
-import math
+from dataclasses import fields, replace
+import json
+import os
 from typing import Sequence
 
 from .core import (
@@ -19,6 +20,53 @@ from .core import (
 )
 
 
+HALT_COT_CONFIG_REPO = os.getenv("HALT_COT_CONFIG_REPO", "yass4/halt-cot")
+
+
+def load_config_from_hub(
+    repo_id: str = HALT_COT_CONFIG_REPO,
+    *,
+    filename: str = "config.json",
+) -> HaltCoTConfig | None:
+    """Fetch the published HALT-CoT defaults from the Hub.
+
+    Downloading ``config.json`` through the Hub also lets Hugging Face count
+    real usage of the method. Any failure (offline, missing repo, malformed
+    file) returns ``None`` so callers fall back to the built-in defaults and a
+    run is never blocked.
+    """
+
+    try:
+        from huggingface_hub import hf_hub_download
+
+        path = hf_hub_download(repo_id=repo_id, filename=filename)
+        payload = json.loads(open(path, encoding="utf-8").read())
+    except Exception:
+        return None
+
+    values = payload.get("halt_cot", payload)
+    if not isinstance(values, dict):
+        return None
+
+    defaults = HaltCoTConfig()
+    tuple_fields = {
+        f.name for f in fields(HaltCoTConfig)
+        if isinstance(getattr(defaults, f.name), tuple)
+    }
+    kwargs = {}
+    for field in fields(HaltCoTConfig):
+        if field.name not in values:
+            continue
+        value = values[field.name]
+        if field.name in tuple_fields and isinstance(value, list):
+            value = tuple(value)
+        kwargs[field.name] = value
+    try:
+        return HaltCoTConfig(**kwargs)
+    except (TypeError, ValueError):
+        return None
+
+
 class TextStopCriteria:
     """Transformers stopping criterion that halts when generated text hits a marker."""
 
@@ -74,6 +122,9 @@ class HaltCoTForCausalLM:
 
         from transformers import AutoModelForCausalLM, AutoTokenizer
 
+        if config is None:
+            config = load_config_from_hub()
+
         tokenizer = AutoTokenizer.from_pretrained(
             model_id,
             trust_remote_code=trust_remote_code,