From 80cb872a84c85fddaa5520f296d94388d732c945 Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Fri, 16 Oct 2026 09:38:30 -0600
Subject: [PATCH] Fall back to an online load when a cached tokenizer/processor
 loads incomplete

---
 toolkit/models/v2/_mixin.py | 34 +++++++++++++++++++++++++++++-----
 1 file changed, 29 insertions(+), 5 deletions(-)

diff --git a/toolkit/models/v2/_mixin.py b/toolkit/models/v2/_mixin.py
index 3446161..9f699fa 100644
--- a/toolkit/models/v2/_mixin.py
+++ b/toolkit/models/v2/_mixin.py
@@ -172,20 +172,36 @@ class OstrisModelMixin:
         return cls._local_first(cls.load_config, path, subfolder=subfolder)
 
     @staticmethod
-    def _local_first(loader, *args, **kwargs):
+    def _local_first(loader, *args, validate=None, **kwargs):
         """Serve fully-cached repos without touching the hub. Online
         from_pretrained does an etag HEAD round-trip per file (seconds of pure
         network latency on every warm load); a local_files_only attempt is
         instant when the cache can't satisfy it, so failing over to the
-        online path costs nothing."""
+        online path costs nothing. ``validate`` rejects a cache-only result
+        that loaded without raising but incomplete (falls over to online)."""
         import os as _os
 
         if _os.environ.get("HF_HUB_OFFLINE") == "1":
             return loader(*args, **kwargs)
         try:
-            return loader(*args, local_files_only=True, **kwargs)
+            result = loader(*args, local_files_only=True, **kwargs)
         except Exception:
             return loader(*args, **kwargs)
+        if validate is not None and not validate(result):
+            try:
+                return loader(*args, **kwargs)
+            except Exception:
+                return result  # no hub: keep what the cache had
+        return result
+
+    @staticmethod
+    def _tokenizer_complete(result) -> bool:
+        # transformers 5.5.3 local_files_only on a partially cached repo (e.g.
+        # config.json only) returns a 1-token tokenizer instead of raising
+        tokenizer = getattr(result, "tokenizer", result)
+        if not hasattr(tokenizer, "__len__"):
+            return True
+        return len(tokenizer) >= 1000
 
     @classmethod
     def aitk_from_config(cls, config):
@@ -803,7 +819,11 @@ class OstrisModelMixin:
         ):
             subfolder = None
         return cls._local_first(
-            AutoTokenizer.from_pretrained, source, subfolder=subfolder or "", **kwargs
+            AutoTokenizer.from_pretrained,
+            source,
+            subfolder=subfolder or "",
+            validate=cls._tokenizer_complete,
+            **kwargs,
         )
 
     @classmethod
@@ -825,7 +845,11 @@ class OstrisModelMixin:
         ):
             subfolder = None
         return cls._local_first(
-            AutoProcessor.from_pretrained, source, subfolder=subfolder or "", **kwargs
+            AutoProcessor.from_pretrained,
+            source,
+            subfolder=subfolder or "",
+            validate=cls._tokenizer_complete,
+            **kwargs,
         )
 
     # ------------------------------------------------------------------