chore: add architectures for config

config.json

@@ -1,4 +1,13 @@
 {
   "model_type": "onnx-base",
-  "model_name_or_path": "model.onnx"
-}
+  "model_path": "model.onnx",
+  "architectures": ["ONNXBaseModel"],
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1"
+  },
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1
+  }
+}

src/demo.py

@@ -0,0 +1,58 @@
+# 1. 首先，你需要定义一个 ONNX 模型配置类，并注册它
+from transformers import AutoConfig, PretrainedConfig, PreTrainedModel, AutoModel
+from transformers.pipelines import PIPELINE_REGISTRY
+
+class ONNXBaseConfig(PretrainedConfig):
+    model_type = "onnx-base"
+
+# 注册配置类
+AutoConfig.register("onnx-base", ONNXBaseConfig)
+
+# 注册模型类
+class ONNXBaseModel(AutoModel):
+    config_class = ONNXBaseConfig
+
+class ONNXBaseModel(PreTrainedModel):
+    config_class = ONNXBaseConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+
+    def forward(self, *args, **kwargs):
+        return self.dummy_param
+
+
+AutoModel.register(ONNXBaseConfig, ONNXBaseModel)
+
+from transformers.pipelines import Pipeline
+
+class ONNXBasePipeline(Pipeline):
+    def __init__(self, model, **kwargs):
+        super().__init__(model=model, **kwargs)
+
+    def _sanitize_parameters(self, **kwargs):
+        return {}, {}, {}
+
+    def preprocess(self, inputs):
+        return inputs
+
+    def _forward(self, model_inputs):
+        return self.model(**model_inputs)
+
+    def postprocess(self, model_outputs):
+        return model_outputs
+
+PIPELINE_REGISTRY.register_pipeline(
+    task="onnx-base",
+    pipeline_class=ONNXBasePipeline
+)
+
+
+from transformers import pipeline
+
+# 使用自定义的 pipeline 任务
+onnx_pipeline = pipeline(task="onnx-base", model="m3/onnx-base")
+
+# 使用 pipeline
+result = onnx_pipeline("Your input data here")
+print(result)

src/init_model.py

@@ -6,7 +6,7 @@ import torch.onnx
 class SimpleModel(nn.Module):
     def __init__(self):
         super(SimpleModel, self).__init__()
-        self.fc = nn.Linear(10, 1)
+        self.fc = nn.Linear(1, 1)
 
     def forward(self, x):
         return self.fc(x)