chore: add empty sample

src/empty.py

@@ -0,0 +1,98 @@
+from transformers import PretrainedConfig, PreTrainedModel, AutoConfig, AutoModel
+from transformers.pipelines import PIPELINE_REGISTRY
+from huggingface_hub import hf_hub_download
+import onnxruntime as ort
+
+import torch
+import os
+
+
+# 1. register AutoConfig
+class ONNXBaseConfig(PretrainedConfig):
+    model_type = 'onnx-base'
+
+
+AutoConfig.register('onnx-base', ONNXBaseConfig)
+
+# 2. register AutoModel
+class ONNXBaseModel(PreTrainedModel):
+    config_class = ONNXBaseConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+
+    def forward(self, input=None, **kwargs):
+        return {}
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
+        return cls(config)
+
+    @property
+    def device(self):
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        return torch.device(device)
+
+
+AutoModel.register(ONNXBaseConfig, ONNXBaseModel)
+
+# 2. register Pipeline
+from transformers.pipelines import Pipeline
+
+class ONNXBasePipeline(Pipeline):
+    def __init__(self, model, **kwargs):
+        super().__init__(model=model, **kwargs)
+        self.device_id = kwargs['device']
+        model_path = hf_hub_download(repo_id='m3/onnx-base', filename='model.onnx', local_files_only=True)
+        self.session = ort.InferenceSession(model_path)
+
+    def __call__(
+            self,
+            inputs: str,
+            **kwargs,
+    ):
+        inputs = {"inputs": inputs}
+        return super().__call__(inputs, **kwargs)
+
+    def _sanitize_parameters(self, **kwargs):
+        return {}, {}, {}
+
+    def preprocess(self, input):
+        return {'input': input}
+
+    def _forward(self, model_input):
+        input = model_input['input']['inputs']
+        outs = self.session.run(None, {'input': input})
+        return input
+
+    def postprocess(self, model_outputs):
+        return model_outputs
+
+
+PIPELINE_REGISTRY.register_pipeline(
+    task='onnx-base',
+    pipeline_class=ONNXBasePipeline,
+    pt_model=ONNXBaseModel,
+    default={"pt": ("m3/onnx-base", "a5e4e8f")},
+)
+
+# 4. show how to use
+from transformers import pipeline
+
+cfg = ONNXBaseConfig(model_path='model.onnx',
+                        id2label={0: 'label_0', 1: 'label_1'},
+                        label2id={0: 'label_1', 1: 'label_0'})
+
+pipe = pipeline(
+    task='onnx-base',
+    batch_size=10,
+    device='cuda',
+)
+
+dummy_input = torch.tensor([[[[1, 2, 3], [4, 5, 6], [7, 8, 9]]]], dtype=torch.float32)
+input_data = dummy_input.numpy()
+result = pipe(
+    inputs=input_data, device='cuda',
+)
+print(result)