Runner with Diaparser Model

[jlunder00]

I am trying to use diaparser (https://github.com/Unipisa/diaparser), to make dependency tree parses from sentences. I want this to be in the form of a runner, so that I can use Bentoml's adaptive batching to make full use of the gpu capabilities of the machine, as a single parse run can take up to 30 seconds. The way the diaparser exposes the capability for you to preform a prediction is as follows:

from diaparser.parsers import Parser

diaparser_parser_object = Parser.load("en_ewt-electra")
result = diaparser_parser_object.predict([['This is a single sentence.']])

I created a runnable class utilizing this predict function, and using the appropriate decorator for a runnable method:

class ParseRunner(bentoml.Runnable):
    SUPPORTED_RESOURCES = ("nvidia.com/gpu",)
    SUPPORTS_CPU_MULTI_THREADING = True

    def __init__(self):
        self.model = self.get_diaparser_model()

    def get_diaparser_model(self):
        return Parser.load("en_ewt-electra").predict

    @bentoml.Runnable.method(batchable=True)
    def parse_sentence(self, text, prob):
        return self.model(text, prob=prob)

Calling this method directly works fine, and behaves as a normal .predict call would. However, when I turn this into a runner:

en_diaparser_runner = bentoml.Runner(ParseRunner, name='en_diaparser_runner', runnable_init_params={})

and try to call .run, I get this error:

File "<stdin>", line 1, in <module>
  File "/home/jlunder/e/e1/lib/python3.8/site-packages/bentoml/_internal/runner/runner.py", line 46, in run
    return self.runner._runner_handle.run_method(  # type: ignore
  File "/home/jlunder/e/e1/lib/python3.8/site-packages/bentoml/_internal/runner/runner_handle/local.py", line 49, in run_method
    return getattr(self._runnable, __bentoml_method.name)(*args, **kwargs)
  File "/home/jlunder/e/e1/lib/python3.8/site-packages/bentoml/_internal/runner/runnable.py", line 139, in method
    return self.func(obj, *args, **kwargs)
  File "/home/jlunder/e/e1/lib/python3.8/site-packages/tree_utils/parse_runner.py", line 36, in parse_sentence
    return self.model(text_preprocessed, prob=prob) if prob is not None else self.model(text_preprocessed)
  File "/home/jlunder/e/e1/lib/python3.8/site-packages/diaparser/parsers/biaffine_dependency.py", line 129, in predict
    return super().predict(**Config().update(locals()))
  File "/home/jlunder/e/e1/lib/python3.8/site-packages/diaparser/parsers/parser.py", line 196, in predict
    preds = self._predict(dataset.loader)
  File "/home/jlunder/e/e1/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/jlunder/e/e1/lib/python3.8/site-packages/diaparser/parsers/biaffine_dependency.py", line 190, in _predict
    for words, feats in progress_bar(loader):
  File "/home/jlunder/e/e1/lib/python3.8/site-packages/tqdm/std.py", line 1183, in __iter__
    for obj in iterable:
  File "/home/jlunder/e/e1/lib/python3.8/site-packages/diaparser/utils/data.py", line 111, in __iter__
    for batch in super().__iter__():
  File "/home/jlunder/e/e1/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 628, in __next__
    data = self._next_data()
  File "/home/jlunder/e/e1/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 671, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "/home/jlunder/e/e1/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 55, in fetch
    if hasattr(self.dataset, "__getitems__") and self.dataset.__getitems__:
  File "/home/jlunder/e/e1/lib/python3.8/site-packages/diaparser/utils/data.py", line 65, in __getattr__
    return [getattr(sentence, name) for sentence in self.sentences]
  File "/home/jlunder/e/e1/lib/python3.8/site-packages/diaparser/utils/data.py", line 65, in <listcomp>
    return [getattr(sentence, name) for sentence in self.sentences]
  File "/home/jlunder/e/e1/lib/python3.8/site-packages/diaparser/utils/transform.py", line 117, in __getattr__
    return self.values[self.maps[name]]
KeyError: '__getitems__'

Which leads me to believe that the object I get from diaparser.parsers.Parser.load is not able to be used as a runner object. Unless I am doing something wrong? Does anyone have any ideas as to how to accomplish this?