Fall back to root tokenizers for ONNX and LiteRT Hub subdirectory models (#1277)

src/python/txtai/vectors/dense/litert.py

@@ -13,7 +13,7 @@ try:
 except ImportError:
     LITERT = False
 
-from ...util import Download, Library
+from ...util import Download, DownloadError, Library
 
 from ..base import Vectors
 
@@ -55,7 +55,16 @@ class LiteRT(Vectors):
         tokenizer = self.config.get("tokenizer")
         if not tokenizer:
             tokenizer = os.path.dirname(path) + "/" + "tokenizer.json"
-            tokenizer = tokenizer if os.path.exists(tokenizer) else Download()(tokenizer)
+            try:
+                tokenizer = tokenizer if os.path.exists(tokenizer) else Download()(tokenizer)
+
+            # Model is in a subdirectory of a HF Hub repo, read the tokenizer from the repo root
+            except DownloadError:
+                # Local models don't have a HF Hub repo root
+                if os.path.exists(path):
+                    raise
+
+                tokenizer = Download()("/".join(path.split("/")[:2]), "tokenizer.json")
 
         # Load tokenizer and model
         tokenizer = Tokenizer.from_file(tokenizer) if os.path.exists(tokenizer) else Tokenizer.from_pretrained(tokenizer)

src/python/txtai/vectors/dense/onnx.py

@@ -14,7 +14,7 @@ try:
 except ImportError:
     ONNX_RUNTIME = False
 
-from ...util import Download, Library
+from ...util import Download, DownloadError, Library
 
 from ..base import Vectors
 
@@ -68,7 +68,7 @@ class ONNX(Vectors):
     def loadtokenizer(self, path):
         """
         Loads the tokenizer for a model. Reads the `tokenizer` configuration option when set,
-        otherwise falls back to a tokenizer.json file stored alongside the model.
+        otherwise falls back to a tokenizer.json file stored alongside the model or at the root of its HF Hub repo.
 
         Args:
             path: model path
@@ -81,7 +81,16 @@ class ONNX(Vectors):
         tokenizer = self.config.get("tokenizer")
         if not tokenizer:
             tokenizer = os.path.dirname(path) + "/" + "tokenizer.json"
-            tokenizer = tokenizer if os.path.exists(tokenizer) else Download()(tokenizer)
+            try:
+                tokenizer = tokenizer if os.path.exists(tokenizer) else Download()(tokenizer)
+
+            # Model is in a subdirectory of a HF Hub repo, read the tokenizer from the repo root
+            except DownloadError:
+                # Local models don't have a HF Hub repo root
+                if os.path.exists(path):
+                    raise
+
+                tokenizer = Download()("/".join(path.split("/")[:2]), "tokenizer.json")
 
         # Local tokenizer file vs a model id resolved through the HF Hub
         return Tokenizer.from_file(tokenizer) if os.path.exists(tokenizer) else Tokenizer.from_pretrained(tokenizer)

test/python/testvectors/testdense/testlitert.py

@@ -3,10 +3,15 @@ LiteRT module tests
 """
 
 import os
+import tempfile
 import unittest
 
+from unittest.mock import patch
+
 import numpy as np
 
+from huggingface_hub import hf_hub_download
+from txtai.util import DownloadError
 from txtai.vectors import VectorsFactory
 
 
@@ -40,3 +45,36 @@ class TestLiteRT(unittest.TestCase):
         # Test shape of serialized embeddings
         with open(stream, "rb") as queue:
             self.assertEqual(np.load(queue).shape, (1, 128))
+
+    def testTokenizerMissing(self):
+        """
+        Test that a local model without a tokenizer raises DownloadError
+        """
+
+        with tempfile.TemporaryDirectory() as directory:
+            path = os.path.join(directory, "model.tflite")
+            with open(path, "wb"):
+                pass
+
+            with self.assertRaises(DownloadError):
+                VectorsFactory.create({"path": path, "gpu": False}, None)
+
+    @patch("huggingface_hub.hf_hub_download")
+    def testTokenizerRepo(self, download):
+        """
+        Test that a tokenizer at the root of a HF Hub repo is found for a model stored in a subdirectory
+        """
+
+        def filedownload(**kwargs):
+            # Serve the model from a subdirectory that doesn't have a tokenizer
+            if kwargs["filename"] == "litert/tokenizer.json":
+                raise FileNotFoundError
+
+            return hf_hub_download(repo_id=kwargs["repo_id"], filename=kwargs["filename"].replace("litert/", ""))
+
+        download.side_effect = filedownload
+
+        model = VectorsFactory.create(
+            {"path": "neuml/bert-hash-nano-embeddings-litert/litert/bert-hash-nano-embeddings-int4.tflite", "gpu": False}, None
+        )
+        self.assertEqual(model.encode(["test"]).shape, (1, 128))

test/python/testvectors/testdense/testonnx.py

@@ -16,6 +16,7 @@ from onnx import helper, TensorProto
 from tokenizers import Tokenizer, models, pre_tokenizers
 
 from txtai.pipeline import HFOnnx
+from txtai.util import DownloadError
 from txtai.vectors import VectorsFactory
 from txtai.vectors.dense.onnx import ONNX
 
@@ -232,6 +233,37 @@ class TestONNXModels(unittest.TestCase):
 
         self.assertEqual(model.encode(["dog"]).shape, (1, 4))
 
+    def testTokenizerMissing(self):
+        """
+        Test that a local model without a tokenizer raises DownloadError
+        """
+
+        os.makedirs(os.path.join(self.directory, "missing-tokenizer"), exist_ok=True)
+        path = self.build(os.path.join("missing-tokenizer", "model.onnx"))
+
+        with self.assertRaises(DownloadError):
+            ONNX({"path": path, "gpu": False}, None, None)
+
+    @patch("huggingface_hub.hf_hub_download")
+    def testTokenizerRepo(self, download):
+        """
+        Test that a tokenizer at the root of a HF Hub repo is found for a model stored in a subdirectory
+        """
+
+        files = {"org/repo/onnx/model.onnx": self.build("repo.onnx"), "org/repo/tokenizer.json": self.tokenizer}
+
+        def filedownload(**kwargs):
+            path = f"{kwargs['repo_id']}/{kwargs['filename']}"
+            if path not in files:
+                raise FileNotFoundError
+
+            return files[path]
+
+        download.side_effect = filedownload
+
+        model = ONNX({"path": "org/repo/onnx/model.onnx", "gpu": False}, None, None)
+        self.assertEqual(model.encode(["dog"]).shape, (1, 4))
+
     def testProviders(self):
         """
         Test provider selection