pipeline

pipe.py

@@ -0,0 +1,33 @@
+import numpy as np
+import torch
+import torch.nn as nn
+
+from transformers import Pipeline
+
+class RegressionPipeline(Pipeline):
+    def _sanitize_parameters(self, **kwargs):
+        preprocess_kwargs = {}
+        if "maybe_arg" in kwargs:
+            preprocess_kwargs["maybe_arg"] = kwargs["maybe_arg"]
+        return preprocess_kwargs, {}, {}
+
+    def preprocess(self, inputs, maybe_arg=2):
+        print(inputs)
+        encoded_corpus = self.tokenizer(text=inputs,
+                                        add_special_tokens=True,
+                                        padding='max_length',
+                                        truncation='longest_first',
+                                        max_length=300,
+                                        return_attention_mask=True)
+        return {"model_input": encoded_corpus}
+
+    def _forward(self, model_inputs):
+        print(model_inputs)
+        # model_inputs == {"model_input": model_input}
+        outputs = self.model(torch.tensor(model_inputs['model_input']['input_ids']).reshape(1, -1).to(torch.int64),
+                             torch.tensor(model_inputs['model_input']['attention_mask']).reshape(1, -1).to(torch.int64))
+        return outputs
+
+    def postprocess(self, model_outputs):
+        print(model_outputs)
+        return model_outputs.numpy()