Adds pipeline model caching in the transform function.
#593
Conversation
| task = json.loads(task) | ||
| args = json.loads(args) | ||
| inputs = json.loads(inputs) | ||
|
|
||
| pipe = transformers.pipeline(**task) | ||
| model = task.get("model") |
There was a problem hiding this comment.
I think there may be different pipelines with the same model, to handle different tasks. We may need to use the full task param for caching. Something like:
if cache:
key = ','.join([str(key) + ':' + str(value) for (key, value) in sorted(task.items())])
if key not in __cache_transformer_by_task:
__cache_transformer_by_task[key] = transformers.pipeline(**task)
pipe = __cache_transformer_by_task[key]
else:
pipe = transformers.pipeline(**task)
There was a problem hiding this comment.
I see, good point. I'll think some more about the parameters of .pipeline() and push an update.
There was a problem hiding this comment.
Okay, I pretty much copied your code verbatim. I also like the idea in there where you have to pass the cache parameter also to USE the cached model. My initial version always used the cached version if it was available whether or not the cache flag was passed in.
|
This looks good to me as an improvement over the status quo, so I'm going to merge it. We're being slightly inconsistent, i.e. every other API caches by default, including the |
This PR adds the ability to run a query that looks like this:
roberta-large-mnliwill be cached in memory to preventtransformers.pipeline()from being called more than once for the same model.By default,
cacheisFALSE.