|
import os |
|
from flash.text import TextClassifier |
|
|
|
|
|
from flash.core.data.data_source import LabelsState |
|
|
|
|
|
class PreTrainedPipeline(): |
|
|
|
def __init__(self, path=""): |
|
self.device = 'cpu' |
|
self.model = TextClassifier.load_from_checkpoint(os.path.join(path, "pytorch_model.bin")) |
|
self.data_pipeline = self.model.build_data_pipeline() |
|
self.labels = self.model._data_pipeline_state._state[LabelsState].labels |
|
self.top_k = 5 |
|
|
|
def __call__(self, inputs): |
|
x = self.data_pipeline._deserializer(inputs) |
|
x = self.data_pipeline.worker_preprocessor('predict')(x) |
|
x = self.model.transfer_batch_to_device(x, self.device, 0) |
|
x = self.data_pipeline.device_preprocessor('predict')(x) |
|
out = self.model.predict_step(x, 0) |
|
proba = out['logits'].softmax(1)[0].tolist() |
|
return [{'score': s, 'label': l} for s, l in sorted(zip(proba, self.labels), reverse=True)[:self.top_k]] |
|
|
