From 18ebcaa735b6e98d9800ab8a37bde2647b310b00 Mon Sep 17 00:00:00 2001 From: Thomas Ma Date: Tue, 30 Aug 2022 17:25:54 -0700 Subject: [PATCH] fixed docstring typo --- opensearch_py_ml/sagemaker_tools.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/opensearch_py_ml/sagemaker_tools.py b/opensearch_py_ml/sagemaker_tools.py index fca13d395..c3034ce9c 100644 --- a/opensearch_py_ml/sagemaker_tools.py +++ b/opensearch_py_ml/sagemaker_tools.py @@ -6,7 +6,7 @@ import numpy as np from opensearch_py_ml import DataFrame -from typing import List, Optional +from typing import List, Optional, Dict, Tuple, Any from math import ceil from sagemaker import RealTimePredictor, Session @@ -21,7 +21,7 @@ def make_sagemaker_prediction(endpoint_name: str, column_order: Optional[List[str]] = None, chunksize: int = None, sort_index: Optional[str] = '_doc' - )-> np.array: + ) -> Tuple[List[Any], Dict[Any, Any]]: """ Make a prediction on an opensearch_py_ml dataframe using a deployed SageMaker model endpoint. @@ -44,7 +44,7 @@ def make_sagemaker_prediction(endpoint_name: str, Returns ---------- - np.array representing the output of the model on input data + list representing the indices, dictionary representing the output of the model on input data """ predictor = RealTimePredictor(endpoint=endpoint_name, sagemaker_session=sagemaker_session, content_type='text/csv') data = data.drop(columns=target_column)