expose functions through pipeline.py

neuralmagic · Jul 17, 2023 · b94937d · b94937d
1 parent deccb3c
commit b94937d
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 4 deletions.
diff --git a/src/deepsparse/pipeline.py b/src/deepsparse/pipeline.py
@@ -26,7 +26,7 @@
 from pydantic import BaseModel, Field
 
 from deepsparse import Context, Engine, MultiModelEngine, Scheduler
-from deepsparse.base_pipeline import BasePipeline
+from deepsparse.base_pipeline import _REGISTERED_PIPELINES, BasePipeline, SupportedTasks
 from deepsparse.benchmark import ORTEngine
 from deepsparse.cpu import cpu_details
 from deepsparse.loggers.base_logger import BaseLogger
@@ -45,6 +45,9 @@
     "ORT_ENGINE",
     "SUPPORTED_PIPELINE_ENGINES",
     "Pipeline",
+    "BasePipeline",
+    "SupportedTasks",
+    "_REGISTERED_PIPELINES",
     "PipelineConfig",
     "question_answering_pipeline",
     "text_classification_pipeline",

diff --git a/src/deepsparse/server/cli.py b/src/deepsparse/server/cli.py
@@ -27,7 +27,7 @@
 import click
 import yaml
 
-from deepsparse.base_pipeline import SupportedTasks
+from deepsparse.pipeline import SupportedTasks
 from deepsparse.server.config import EndpointConfig, ServerConfig
 from deepsparse.server.server import start_server
 

diff --git a/tests/deepsparse/pipelines/test_dynamic_import.py b/tests/deepsparse/pipelines/test_dynamic_import.py
@@ -15,8 +15,7 @@
 import os
 
 import pytest
-from deepsparse.base_pipeline import _REGISTERED_PIPELINES
-from deepsparse.pipeline import Pipeline
+from deepsparse.pipeline import _REGISTERED_PIPELINES, Pipeline
 from deepsparse.tasks import _split_dir_and_name, dynamic_import_task