fix tests

Signed-off-by: Raphael Glon <[email protected]>

Author: oOraph

src/huggingface_inference_toolkit/webservice_starlette.py

@@ -27,10 +27,6 @@
 from huggingface_inference_toolkit.serialization.base import ContentType
 from huggingface_inference_toolkit.serialization.json_utils import Jsoner
 from huggingface_inference_toolkit.utils import convert_params_to_int_or_bool
-
-#     _load_repository_from_hf,
-#     convert_params_to_int_or_bool,
-# )
 from huggingface_inference_toolkit.vertex_ai_utils import _load_repository_from_gcs
 
 INFERENCE_HANDLERS = {}

tests/integ/conftest.py

@@ -9,7 +9,7 @@
 import tenacity
 from transformers.testing_utils import _run_slow_tests
 
-from huggingface_inference_toolkit.utils import _load_repository_from_hf
+from huggingface_inference_toolkit.heavy_utils import load_repository_from_hf
 from tests.integ.config import task2model
 
 HF_HUB_CACHE = os.environ.get("HF_HUB_CACHE", "/home/ubuntu/.cache/huggingface/hub")
@@ -124,7 +124,7 @@ def local_container(device, task, repository_id, framework):
         object_id = model.replace("/", "--")
         model_dir = f"{HF_HUB_CACHE}/{object_id}"
 
-        _storage_dir = _load_repository_from_hf(
+        _storage_dir = load_repository_from_hf(
             repository_id=model, target_dir=model_dir
         )
 

tests/integ/helpers.py

@@ -10,7 +10,7 @@
 from docker import DockerClient
 from transformers.testing_utils import _run_slow_tests, require_tf, require_torch
 
-from huggingface_inference_toolkit.utils import _load_repository_from_hf
+from huggingface_inference_toolkit.heavy_utils import load_repository_from_hf
 from tests.integ.config import task2input, task2model, task2output, task2validation
 
 IS_GPU = _run_slow_tests
@@ -207,7 +207,7 @@ def test_pt_container_local_model(task: str) -> None:
     make_sure_other_containers_are_stopped(client, container_name)
     with tempfile.TemporaryDirectory() as tmpdirname:
         # https://github.com/huggingface/infinity/blob/test-ovh/test/integ/utils.py
-        _load_repository_from_hf(model, tmpdirname, framework="pytorch")
+        load_repository_from_hf(model, tmpdirname, framework="pytorch")
         container = client.containers.run(
             container_image,
             name=container_name,
@@ -238,7 +238,7 @@ def test_pt_container_custom_handler(repository_id) -> None:
     make_sure_other_containers_are_stopped(client, container_name)
     with tempfile.TemporaryDirectory() as tmpdirname:
         # https://github.com/huggingface/infinity/blob/test-ovh/test/integ/utils.py
-        _storage_dir = _load_repository_from_hf(repository_id, tmpdirname)
+        _storage_dir = load_repository_from_hf(repository_id, tmpdirname)
         container = client.containers.run(
             container_image,
             name=container_name,
@@ -275,7 +275,7 @@ def test_pt_container_legacy_custom_pipeline(repository_id: str) -> None:
     make_sure_other_containers_are_stopped(client, container_name)
     with tempfile.TemporaryDirectory() as tmpdirname:
         # https://github.com/huggingface/infinity/blob/test-ovh/test/integ/utils.py
-        _storage_dir = _load_repository_from_hf(repository_id, tmpdirname)
+        _storage_dir = load_repository_from_hf(repository_id, tmpdirname)
         container = client.containers.run(
             container_image,
             name=container_name,
@@ -393,7 +393,7 @@ def test_tf_container_local_model(task) -> None:
     make_sure_other_containers_are_stopped(client, container_name)
     with tempfile.TemporaryDirectory() as tmpdirname:
         # https://github.com/huggingface/infinity/blob/test-ovh/test/integ/utils.py
-        _storage_dir = _load_repository_from_hf(model, tmpdirname, framework=framework)
+        _storage_dir = load_repository_from_hf(model, tmpdirname, framework=framework)
         container = client.containers.run(
             container_image,
             name=container_name,
@@ -421,7 +421,7 @@ def test_tf_container_local_model(task) -> None:
 #     make_sure_other_containers_are_stopped(client, container_name)
 #     with tempfile.TemporaryDirectory() as tmpdirname:
 #         # https://github.com/huggingface/infinity/blob/test-ovh/test/integ/utils.py
-#         storage_dir = _load_repository_from_hf("philschmid/custom-pipeline-text-classification", tmpdirname)
+#         storage_dir = load_repository_from_hf("philschmid/custom-pipeline-text-classification", tmpdirname)
 #         container = client.containers.run(
 #             container_image,
 #             name=container_name,

tests/unit/test_diffusers.py

@@ -5,14 +5,14 @@
 from transformers.testing_utils import require_torch, slow
 
 from huggingface_inference_toolkit.diffusers_utils import IEAutoPipelineForText2Image
-from huggingface_inference_toolkit.utils import _load_repository_from_hf, get_pipeline
+from huggingface_inference_toolkit.heavy_utils import get_pipeline, load_repository_from_hf
 
 logging.basicConfig(level="DEBUG")
 
 @require_torch
 def test_get_diffusers_pipeline():
     with tempfile.TemporaryDirectory() as tmpdirname:
-        storage_dir = _load_repository_from_hf(
+        storage_dir = load_repository_from_hf(
             "echarlaix/tiny-random-stable-diffusion-xl",
             tmpdirname,
             framework="pytorch"
@@ -25,7 +25,7 @@ def test_get_diffusers_pipeline():
 @require_torch
 def test_pipe_on_gpu():
     with tempfile.TemporaryDirectory() as tmpdirname:
-        storage_dir = _load_repository_from_hf(
+        storage_dir = load_repository_from_hf(
             "echarlaix/tiny-random-stable-diffusion-xl",
             tmpdirname,
             framework="pytorch"
@@ -41,7 +41,7 @@ def test_pipe_on_gpu():
 @require_torch
 def test_text_to_image_task():
     with tempfile.TemporaryDirectory() as tmpdirname:
-        storage_dir = _load_repository_from_hf(
+        storage_dir = load_repository_from_hf(
             "echarlaix/tiny-random-stable-diffusion-xl",
             tmpdirname,
             framework="pytorch"

tests/unit/test_handler.py

@@ -8,9 +8,9 @@
     HuggingFaceHandler,
     get_inference_handler_either_custom_or_default_handler,
 )
-from huggingface_inference_toolkit.utils import (
+from huggingface_inference_toolkit.heavy_utils import (
     _is_gpu_available,
-    _load_repository_from_hf,
+    load_repository_from_hf,
 )
 
 TASK = "text-classification"
@@ -29,7 +29,7 @@ def test_pt_get_device() -> None:
 
     with tempfile.TemporaryDirectory() as tmpdirname:
         # https://github.com/huggingface/infinity/blob/test-ovh/test/integ/utils.py
-        storage_dir = _load_repository_from_hf(MODEL, tmpdirname, framework="pytorch")
+        storage_dir = load_repository_from_hf(MODEL, tmpdirname, framework="pytorch")
         h = HuggingFaceHandler(model_dir=str(storage_dir), task=TASK)
         if torch.cuda.is_available():
             assert h.pipeline.model.device == torch.device(type="cuda", index=0)
@@ -41,7 +41,7 @@ def test_pt_get_device() -> None:
 def test_pt_predict_call(input_data: Dict[str, str]) -> None:
     with tempfile.TemporaryDirectory() as tmpdirname:
         # https://github.com/huggingface/infinity/blob/test-ovh/test/integ/utils.py
-        storage_dir = _load_repository_from_hf(MODEL, tmpdirname, framework="pytorch")
+        storage_dir = load_repository_from_hf(MODEL, tmpdirname, framework="pytorch")
         h = HuggingFaceHandler(model_dir=str(storage_dir), task=TASK)
 
         prediction = h(input_data)
@@ -52,7 +52,7 @@ def test_pt_predict_call(input_data: Dict[str, str]) -> None:
 @require_torch
 def test_pt_custom_pipeline(input_data: Dict[str, str]) -> None:
     with tempfile.TemporaryDirectory() as tmpdirname:
-        storage_dir = _load_repository_from_hf(
+        storage_dir = load_repository_from_hf(
             "philschmid/custom-pipeline-text-classification",
             tmpdirname,
             framework="pytorch",
@@ -64,7 +64,7 @@ def test_pt_custom_pipeline(input_data: Dict[str, str]) -> None:
 @require_torch
 def test_pt_sentence_transformers_pipeline(input_data: Dict[str, str]) -> None:
     with tempfile.TemporaryDirectory() as tmpdirname:
-        storage_dir = _load_repository_from_hf(
+        storage_dir = load_repository_from_hf(
             "sentence-transformers/all-MiniLM-L6-v2", tmpdirname, framework="pytorch"
         )
         h = get_inference_handler_either_custom_or_default_handler(str(storage_dir), task="sentence-embeddings")
@@ -76,7 +76,7 @@ def test_pt_sentence_transformers_pipeline(input_data: Dict[str, str]) -> None:
 def test_tf_get_device():
     with tempfile.TemporaryDirectory() as tmpdirname:
         # https://github.com/huggingface/infinity/blob/test-ovh/test/integ/utils.py
-        storage_dir = _load_repository_from_hf(MODEL, tmpdirname, framework="tensorflow")
+        storage_dir = load_repository_from_hf(MODEL, tmpdirname, framework="tensorflow")
         h = HuggingFaceHandler(model_dir=str(storage_dir), task=TASK)
         if _is_gpu_available():
             assert h.pipeline.device == 0
@@ -88,7 +88,7 @@ def test_tf_get_device():
 def test_tf_predict_call(input_data: Dict[str, str]) -> None:
     with tempfile.TemporaryDirectory() as tmpdirname:
         # https://github.com/huggingface/infinity/blob/test-ovh/test/integ/utils.py
-        storage_dir = _load_repository_from_hf(MODEL, tmpdirname, framework="tensorflow")
+        storage_dir = load_repository_from_hf(MODEL, tmpdirname, framework="tensorflow")
         handler = HuggingFaceHandler(model_dir=str(storage_dir), task=TASK, framework="tf")
 
         prediction = handler(input_data)
@@ -99,7 +99,7 @@ def test_tf_predict_call(input_data: Dict[str, str]) -> None:
 @require_tf
 def test_tf_custom_pipeline(input_data: Dict[str, str]) -> None:
     with tempfile.TemporaryDirectory() as tmpdirname:
-        storage_dir = _load_repository_from_hf(
+        storage_dir = load_repository_from_hf(
             "philschmid/custom-pipeline-text-classification",
             tmpdirname,
             framework="tensorflow",
@@ -112,7 +112,7 @@ def test_tf_custom_pipeline(input_data: Dict[str, str]) -> None:
 def test_tf_sentence_transformers_pipeline():
     # TODO should fail! because TF is not supported yet
     with tempfile.TemporaryDirectory() as tmpdirname:
-        storage_dir = _load_repository_from_hf(
+        storage_dir = load_repository_from_hf(
             "sentence-transformers/all-MiniLM-L6-v2", tmpdirname, framework="tensorflow"
         )
         with pytest.raises(Exception) as _exc_info:

tests/unit/test_optimum_utils.py

@@ -4,12 +4,12 @@
 import pytest
 from transformers.testing_utils import require_torch
 
+from huggingface_inference_toolkit.heavy_utils import load_repository_from_hf
 from huggingface_inference_toolkit.optimum_utils import (
     get_input_shapes,
     get_optimum_neuron_pipeline,
     is_optimum_neuron_available,
 )
-from huggingface_inference_toolkit.utils import _load_repository_from_hf
 
 require_inferentia = pytest.mark.skipif(
     not is_optimum_neuron_available(),
@@ -34,7 +34,7 @@ def test_not_supported_task():
 @require_inferentia
 def test_get_input_shapes_from_file():
     with tempfile.TemporaryDirectory() as tmpdirname:
-        storage_folder = _load_repository_from_hf(
+        storage_folder = load_repository_from_hf(
             repository_id=REMOTE_CONVERTED_MODEL,
             target_dir=tmpdirname,
         )
@@ -49,7 +49,7 @@ def test_get_input_shapes_from_env():
     os.environ["HF_OPTIMUM_BATCH_SIZE"] = "4"
     os.environ["HF_OPTIMUM_SEQUENCE_LENGTH"] = "32"
     with tempfile.TemporaryDirectory() as tmpdirname:
-        storage_folder = _load_repository_from_hf(
+        storage_folder = load_repository_from_hf(
             repository_id=REMOTE_NOT_CONVERTED_MODEL,
             target_dir=tmpdirname,
         )
@@ -77,7 +77,7 @@ def test_get_optimum_neuron_pipeline_from_converted_model():
 def test_get_optimum_neuron_pipeline_from_non_converted_model():
     os.environ["HF_OPTIMUM_SEQUENCE_LENGTH"] = "32"
     with tempfile.TemporaryDirectory() as tmpdirname:
-        storage_folder = _load_repository_from_hf(
+        storage_folder = load_repository_from_hf(
             repository_id=REMOTE_NOT_CONVERTED_MODEL,
             target_dir=tmpdirname,
         )

tests/unit/test_sentence_transformers.py

@@ -3,28 +3,28 @@
 import pytest
 from transformers.testing_utils import require_torch
 
+from huggingface_inference_toolkit.heavy_utils import (
+    get_pipeline,
+    load_repository_from_hf,
+)
 from huggingface_inference_toolkit.sentence_transformers_utils import (
     SentenceEmbeddingPipeline,
     get_sentence_transformers_pipeline,
 )
-from huggingface_inference_toolkit.utils import (
-    _load_repository_from_hf,
-    get_pipeline,
-)
 
 
 @require_torch
 def test_get_sentence_transformers_pipeline():
     with tempfile.TemporaryDirectory() as tmpdirname:
-        storage_dir = _load_repository_from_hf("sentence-transformers/all-MiniLM-L6-v2", tmpdirname)
+        storage_dir = load_repository_from_hf("sentence-transformers/all-MiniLM-L6-v2", tmpdirname)
         pipe = get_pipeline("sentence-embeddings", storage_dir.as_posix())
         assert isinstance(pipe, SentenceEmbeddingPipeline)
 
 
 @require_torch
 def test_sentence_embedding_task():
     with tempfile.TemporaryDirectory() as tmpdirname:
-        storage_dir = _load_repository_from_hf("sentence-transformers/all-MiniLM-L6-v2", tmpdirname)
+        storage_dir = load_repository_from_hf("sentence-transformers/all-MiniLM-L6-v2", tmpdirname)
         pipe = get_sentence_transformers_pipeline("sentence-embeddings", storage_dir.as_posix())
         res = pipe(sentences="Lets create an embedding")
         assert isinstance(res["embeddings"], list)
@@ -36,7 +36,7 @@ def test_sentence_embedding_task():
 @require_torch
 def test_sentence_similarity():
     with tempfile.TemporaryDirectory() as tmpdirname:
-        storage_dir = _load_repository_from_hf("sentence-transformers/all-MiniLM-L6-v2", tmpdirname)
+        storage_dir = load_repository_from_hf("sentence-transformers/all-MiniLM-L6-v2", tmpdirname)
         pipe = get_sentence_transformers_pipeline("sentence-similarity", storage_dir.as_posix())
         res = pipe(source_sentence="Lets create an embedding", sentences=["Lets create an embedding"])
         assert isinstance(res["similarities"], list)
@@ -45,7 +45,7 @@ def test_sentence_similarity():
 @require_torch
 def test_sentence_ranking():
     with tempfile.TemporaryDirectory() as tmpdirname:
-        storage_dir = _load_repository_from_hf("cross-encoder/ms-marco-MiniLM-L-6-v2", tmpdirname)
+        storage_dir = load_repository_from_hf("cross-encoder/ms-marco-MiniLM-L-6-v2", tmpdirname)
         pipe = get_sentence_transformers_pipeline("sentence-ranking", storage_dir.as_posix())
         res = pipe(
             sentences=[
@@ -61,7 +61,7 @@ def test_sentence_ranking():
 @require_torch
 def test_sentence_ranking_tei():
     with tempfile.TemporaryDirectory() as tmpdirname:
-        storage_dir = _load_repository_from_hf("cross-encoder/ms-marco-MiniLM-L-6-v2", tmpdirname, framework="pytorch")
+        storage_dir = load_repository_from_hf("cross-encoder/ms-marco-MiniLM-L-6-v2", tmpdirname, framework="pytorch")
         pipe = get_sentence_transformers_pipeline("sentence-ranking", storage_dir.as_posix())
         res = pipe(
             query="Lets create an embedding",
@@ -82,7 +82,7 @@ def test_sentence_ranking_tei():
 @require_torch
 def test_sentence_ranking_validation_errors():
     with tempfile.TemporaryDirectory() as tmpdirname:
-        storage_dir = _load_repository_from_hf("cross-encoder/ms-marco-MiniLM-L-6-v2", tmpdirname, framework="pytorch")
+        storage_dir = load_repository_from_hf("cross-encoder/ms-marco-MiniLM-L-6-v2", tmpdirname, framework="pytorch")
         pipe = get_sentence_transformers_pipeline("sentence-ranking", storage_dir.as_posix())
 
         with pytest.raises(