chore: add dummpy model

.gitignore

@@ -1 +1,2 @@
 ./custom_model
+.idea

model.onnx

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

src/{empty.py → dummy.py}

@@ -1,51 +1,45 @@
 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):
+# 1. create auto config
+class DummyConfig(PretrainedConfig):
+    model_type = 'dummy'
+
+# 2. create model
+class DummyModel(PreTrainedModel):
+    config_class = DummyConfig
+    def __init__(self, model: str, model_path: str):
+        is_local = os.path.isdir(model)
+        if is_local:
+            base_path = model
+            model_path = os.path.join(base_path, model_path)
+        else:
+            model_path = hf_hub_download(repo_id=model, filename=model_path)
+            base_path = os.path.dirname(model_path)
+        config = DummyConfig(base_path=base_path, model_path=model_path)
         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):
+class DummyPipeline(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)
+        self.model_path = self.model.config.model_path
+        self.base_path = self.model.config.base_path
+
 
     def __call__(
             self,
@@ -62,30 +56,27 @@ class ONNXBasePipeline(Pipeline):
         return {'input': input}
 
     def _forward(self, model_input):
-        input = model_input['input']['inputs']
-        outs = self.session.run(None, {'input': input})
-        return input
+        return {'data': 'dummy',
+                'device_id': self.device_id,
+                'base_path': self.base_path,
+                'model_path': self.model_path
+                }
 
     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")},
+    task='dummy-task',
+    pipeline_class=DummyPipeline,
+    pt_model=DummyModel,
 )
 
 # 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',
+    model=DummyModel("Ultralytics/YOLOv8", 'yolov8m.pt'),
+    task='dummy-task',
     batch_size=10,
     device='cuda',
 )