more fixes (#1459)

Author: dsikka

src/deepsparse/transformers/helpers.py

@@ -68,7 +68,7 @@ def setup_transformers_pipeline(
     :param engine_kwargs: The kwargs to pass to the engine
     :return The model path, config, tokenizer, and engine kwargs
     """
-    model_path, config, tokenizer = setup_onnx_file_path(model_path, sequence_length)
+    model_path, config, tokenizer = fetch_onnx_file_path(model_path, sequence_length)
 
     tokenizer.padding_side = tokenizer_padding_side
     if not tokenizer.pad_token:
@@ -86,7 +86,7 @@ def setup_transformers_pipeline(
     return model_path, config, tokenizer, engine_kwargs
 
 
-def setup_onnx_file_path(
+def fetch_onnx_file_path(
     model_path: str,
     sequence_length: int,
     task: Optional[str] = None,
@@ -135,8 +135,6 @@ def get_deployment_path(model_path: str) -> Tuple[str, str]:
     for running the transformers model in the deepsparse pipeline
 
     :param model_path: path to model directory, sparsezoo stub, or ONNX file
-    :param onnx_model_name: name of the ONNX file to look for in the deployment
-        directory. Defaults to MODEL_ONNX_NAME
     :return: path to the deployment directory and path to the ONNX file inside
         the deployment directory
     """

src/deepsparse/v2/pipeline.py

@@ -54,8 +54,8 @@ def __init__(
         ops: Union[Dict[str, Operator], List[Operator]],
         router: Router,
         schedulers: List[OperatorScheduler],
-        continuous_batching_scheduler: ContinuousBatchingScheduler,
-        pipeline_state: PipelineState = None,
+        continuous_batching_scheduler: Optional[ContinuousBatchingScheduler] = None,
+        pipeline_state: Optional[PipelineState] = None,
     ):
 
         self.ops = ops
@@ -277,7 +277,7 @@ def run(
                 if operator_output is None:
                     raise ValueError(
                         f"{self.router.SPLIT_ROUTE} should appear after "
-                        f"{self.ROUTER.START_ROUTE}"
+                        f"{self.router.START_ROUTE}"
                     )
 
                 operator_output = asyncio.run(

tests/deepsparse/v2/schedulers/test_continuous_batching_scheduler.py

@@ -29,7 +29,6 @@ def test_continuous_batching_executor_thread():
     # mobilenet model with batch_size=2
     engine_operator = EngineOperator(
         "zoo:mobilenet_v2-1.0-imagenet-base",
-        batch_size=1,
     )
 
     scheduler.add_engine_operator(engine_operator, [1])

tests/deepsparse/v2/schedulers/utils/test_continuous_batching_executor.py

@@ -26,14 +26,14 @@
 
 def test_continuous_batching_executor_thread():
     # mobilenet model with batch_size=2
-    engine_operator = EngineOperator("zoo:mobilenet_v2-1.0-imagenet-base", batch_size=2)
+    engine_operator = EngineOperator("zoo:mobilenet_v2-1.0-imagenet-base")
 
     # create queues object and add operator
     queues = ContinuousBatchingQueues()
-    queues.add_queue(engine_operator, batch_sizes=[2])
+    queues.add_queue(engine_operator, batch_sizes=[1])
 
     # create engine map
-    operators_to_engines = {engine_operator: {2: engine_operator.engine}}
+    operators_to_engines = {engine_operator: {1: engine_operator.engine}}
 
     worker_thread = ContinuousBatchingExecutorThread(queues, operators_to_engines)
 

tests/deepsparse/v2/test_image_classification.py

@@ -34,6 +34,10 @@ def test_image_classification(get_images):
         "zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenet/pruned95-none"
     )
     pipeline = ImageClassificationPipeline(model_path=model_path)
-    output = pipeline(ImageClassificationInput(images=get_images))
-    assert output.labels == [[207], [670]]
-    assert numpy.allclose(output.scores, [[21.85], [17.33]], atol=0.01)
+    ground_truth = [[207], [670]]
+    scores = [[21.85], [17.33]]
+
+    for i in range(len(get_images)):
+        output = pipeline(ImageClassificationInput(images=get_images[i]))
+        assert output.labels == ground_truth[i]
+        assert numpy.allclose(output.scores, scores[i], atol=0.01)