feat: stable releases

model.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:535fb5c7cda692ca2caad3c6d15260f7ebe398f7b16556d038b49e34ced0a3e2
+oid sha256:4e25f635133a806a543cc19b1fae70ae4ade6a30d2770a58fae8a69834b5428e
 size 235

src/create_model.py

@@ -0,0 +1,57 @@
+from transformers import PretrainedConfig, PreTrainedModel
+import torch
+import torch.nn as nn
+
+
+class ONNXBaseConfig(PretrainedConfig):
+    model_type = "onnx-base"
+
+    def __init__(self, model_path=None, **kwargs):
+        self.model_path = model_path
+        super().__init__(**kwargs)
+
+
+model_directory = './new_model'
+
+config = ONNXBaseConfig(model_path='model.onnx')
+config.save_pretrained(save_directory=model_directory)
+
+
+class ONNXBaseModel(PreTrainedModel):
+    config_class = ONNXBaseConfig
+    def __init__(self, config):
+        super().__init__(config)
+        self.dummy_param = nn.Parameter(torch.zeros(0))
+
+    def forward(self, inputs):
+        return torch.zeros_like(inputs)
+
+    def save_pretrained(self, save_directory: str, **kwargs):
+        super().save_pretrained(save_directory=save_directory, **kwargs)
+        onnx_file_path = save_directory + '/model.onnx'
+        dummy_input = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32)
+        torch.onnx.export(self, dummy_input, onnx_file_path,
+                          input_names=['input'], output_names=['output'],
+                          dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}})
+
+
+# Initialize model
+model = ONNXBaseModel(config)
+# Save model
+model.save_pretrained(save_directory=model_directory)
+
+model = model.from_pretrained(model_directory)
+
+# Test model
+dummy_input = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32)
+output_tensor = model(dummy_input)
+print(output_tensor)
+
+# Test the onnx model
+onnx_file_path = model_directory + '/model.onnx'
+import onnx
+import onnxruntime as ort
+
+ort_session = ort.InferenceSession(onnx_file_path)
+outputs = ort_session.run(None, {'input': dummy_input.numpy()})
+print("Model output:", outputs)

src/demo.py

@@ -7,14 +7,9 @@ pipe = pipeline(
     task='onnx-base',
     model='m3/onnx-base',
     batch_size=10,
-    device='cuda',
+    device='cpu',
 )
-
-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',
-)
-
+input = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32).numpy()
+result = pipe(input)
 print(result)
 

src/init_model.py

@@ -1,34 +0,0 @@
-import torch
-import os
-import torch.nn as nn
-from pipeline import ONNXBaseConfig, ONNXBaseModel
-
-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'}})
-

src/init_onnx.py

@@ -1,37 +0,0 @@
-import torch
-import torch.nn as nn
-
-class BaseModel(nn.Module):
-    def __init__(self):
-        super(BaseModel, self).__init__()
-
-    def forward(self, 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)
-
-safetensors_file_path = "model.safetensors"
-
-from safetensors.torch import save_file
-save_file(model.state_dict(), 'model.safetensors')
-
-import torch.onnx
-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)
-

src/pipeline.py

@@ -26,6 +26,14 @@ class ONNXBaseModel(PreTrainedModel):
         outs = self.session.run(None, {'input': input})
         return outs
 
+    def save_pretrained(self, save_directory: str, **kwargs):
+        super().save_pretrained(save_directory=save_directory, **kwargs)
+        onnx_file_path = save_directory + '/model.onnx'
+        dummy_input = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32)
+        torch.onnx.export(self, dummy_input, onnx_file_path,
+                          input_names=['input'], output_names=['output'],
+                          dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}})
+
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
         config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
@@ -45,6 +53,7 @@ class ONNXBaseModel(PreTrainedModel):
         device = 'cuda' if torch.cuda.is_available() else 'cpu'
         return torch.device(device)
 
+
 AutoModel.register(ONNXBaseConfig, ONNXBaseModel)
 
 # 2. register Pipeline