|
from typing import Dict, Any |
|
from deepsparse import Pipeline |
|
from time import perf_counter |
|
|
|
class EndpointHandler: |
|
|
|
def __init__(self, path=""): |
|
|
|
self.pipeline = Pipeline.create(task="text-classification", model_path=path, scheduler="sync") |
|
|
|
def __call__(self, data: Dict[str, Any]) -> Dict[str, str]: |
|
""" |
|
Args: |
|
data (:obj:): prediction input text |
|
""" |
|
inputs = data.pop("inputs", data) |
|
|
|
start = perf_counter() |
|
prediction = self.pipeline(inputs) |
|
end = perf_counter() |
|
latency = end - start |
|
|
|
return { |
|
"labels": prediction.labels, |
|
"scores": prediction.scores, |
|
"latency (secs.)": latency |
|
} |