init

README.md

@@ -0,0 +1,7 @@
+---
+tags:
+- text-classification
+library_name: generic
+---
+
+# Test

pipeline.py

@@ -0,0 +1,24 @@
+import os
+from flash.text import TextClassifier
+
+# ⚠️ You need this to access the state key
+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]]

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:783e24a9d32146f692d994e4bd807e54da084a0ab76743584b32da01e1c93dda
+size 165569317

requirements.txt

@@ -0,0 +1 @@
+lightning-flash[text]