feat: add onnx model

.gitignore

@@ -0,0 +1 @@
+./custom_model

src/demo.py

@@ -1,58 +1,124 @@
-# 1. 首先，你需要定义一个 ONNX 模型配置类，并注册它
-from transformers import AutoConfig, PretrainedConfig, PreTrainedModel, AutoModel
+from transformers import PretrainedConfig, PreTrainedModel, AutoConfig, AutoModel, modeling_utils
 from transformers.pipelines import PIPELINE_REGISTRY
+from huggingface_hub import hf_hub_download
 
-class ONNXBaseConfig(PretrainedConfig):
-    model_type = "onnx-base"
+import onnxruntime as ort
+import torch
+import os
+import torch.nn as nn
 
-# 注册配置类
-AutoConfig.register("onnx-base", ONNXBaseConfig)
+# 1. register AutoConfig
+class ONNXBaseConfig(PretrainedConfig):
+    model_type = 'onnx-base'
 
-# 注册模型类
-class ONNXBaseModel(AutoModel):
-    config_class = ONNXBaseConfig
+AutoConfig.register('onnx-base', ONNXBaseConfig)
 
+# 2. register AutoModel
 class ONNXBaseModel(PreTrainedModel):
     config_class = ONNXBaseConfig
-
-    def __init__(self, config):
+    def __init__(self, config, base_path=None):
         super().__init__(config)
-
-    def forward(self, *args, **kwargs):
-        return self.dummy_param
-
+        if base_path:
+            model_path = base_path + '/' + config.model_path
+            if os.path.exists(model_path):
+                self.session = ort.InferenceSession(model_path)
+
+    def forward(self, input=None, **kwargs):
+        outs = self.session.run(None, {'input': input})
+        return outs
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
+        is_local = os.path.isdir(pretrained_model_name_or_path)
+        if is_local:
+            base_path = pretrained_model_name_or_path
+        else:
+            config_path = hf_hub_download(repo_id=pretrained_model_name_or_path, filename='config.json')
+            base_path = os.path.dirname(config_path)
+            hf_hub_download(repo_id=pretrained_model_name_or_path, filename=config.model_path)
+        return cls(config, base_path=base_path)
+
+    @property
+    def device(self):
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        return torch.device(device)
 
 AutoModel.register(ONNXBaseConfig, ONNXBaseModel)
 
+# option: save config to path
+local_model_path = './custom_model'
+config = ONNXBaseConfig(model_path='model.onnx',
+                        id2label={0: 'label_0', 1: 'label_1'},
+                        label2id={0: 'label_1', 1: 'label_0'})
+model = ONNXBaseModel(config, base_path='./custom_mode')
+config.save_pretrained(local_model_path)
+# make sure have model_type
+import json
+config_path = local_model_path +  '/config.json'
+with open(config_path, 'r') as f:
+    config_data = json.load(f)
+config_data['model_type'] = 'onnx-base'
+del config_data['transformers_version']
+with open(config_path, 'w') as f:
+    json.dump(config_data, f, indent=2)
+
+# save onnx
+dummy_input = torch.tensor([[[[1, 2, 3], [4, 5, 6], [7, 8, 9]]]], dtype=torch.float32)
+onnx_file_path = './custom_model' + '/' +  'model.onnx'
+class ZeroModel(nn.Module):
+    def __init__(self):
+        super(ZeroModel, self).__init__()
+    def forward(self, x):
+        return torch.zeros_like(x)
+zero_model = ZeroModel()
+torch.onnx.export(zero_model, dummy_input, onnx_file_path,
+                  input_names=['input'], output_names=['output'],
+                  dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}})
+
+
+# 2. register Pipeline
 from transformers.pipelines import Pipeline
 
 class ONNXBasePipeline(Pipeline):
     def __init__(self, model, **kwargs):
+        self.device_id = kwargs['device']
         super().__init__(model=model, **kwargs)
 
     def _sanitize_parameters(self, **kwargs):
         return {}, {}, {}
 
-    def preprocess(self, inputs):
-        return inputs
+    def preprocess(self, input):
+        return {'input': input}
 
-    def _forward(self, model_inputs):
-        return self.model(**model_inputs)
+    def _forward(self, model_input):
+        with torch.no_grad():
+            outputs = self.model(**model_input)
+        return outputs
 
     def postprocess(self, model_outputs):
         return model_outputs
 
 PIPELINE_REGISTRY.register_pipeline(
-    task="onnx-base",
-    pipeline_class=ONNXBasePipeline
+    task='onnx-base',
+    pipeline_class=ONNXBasePipeline,
+    pt_model=ONNXBaseModel
 )
 
-
+# 4. show how to use
 from transformers import pipeline
 
-# 使用自定义的 pipeline 任务
-onnx_pipeline = pipeline(task="onnx-base", model="m3/onnx-base")
+pipe = pipeline(
+    task='onnx-base',
+    model='m3/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)
 
-# 使用 pipeline
-result = onnx_pipeline("Your input data here")
-print(result)

src/init_model.py

@@ -1,18 +1,31 @@
 import torch
 import torch.nn as nn
 import torch.onnx
-
-# Define a simple model
-class SimpleModel(nn.Module):
+class BaseModel(nn.Module):
     def __init__(self):
-        super(SimpleModel, self).__init__()
-        self.fc = nn.Linear(1, 1)
+        super(BaseModel, self).__init__()
 
     def forward(self, x):
-        return self.fc(x)
+        return torch.zeros_like(x)
+
+# create a model
+model = BaseModel()
+
+dummy_input = torch.tensor([[[[1, 2, 3], [4, 5, 6], [7, 8, 9]]]], dtype=torch.float32)
+
+onnx_file_path = "model.onnx"
+torch.onnx.export(model, dummy_input, onnx_file_path,
+                  input_names=['input'], output_names=['output'],
+                  dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}})
+
+print(f"Model has been exported to {onnx_file_path}")
+
+import onnx
+import onnxruntime as ort
+onnx_model = onnx.load(onnx_file_path)
+onnx.checker.check_model(onnx_model)
+ort_session = ort.InferenceSession(onnx_file_path)
+input_data = dummy_input.numpy()
+outputs = ort_session.run(None, {'input': input_data})
+print("Model output:", outputs)
 
-# Instantiate and export the model
-model = SimpleModel()
-dummy_input = torch.randn(1, 10)
-onnx_path = "../model.onnx"
-torch.onnx.export(model, dummy_input, onnx_path, input_names=['input'], output_names=['output'])