Delete onnx_export.py

onnx_export.py

@@ -1,62 +0,0 @@
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import torch
-import onnx
-import onnxruntime as ort
-import numpy as np
-
-
-model_dir = './finetune/gte-base-custom-matryoshka'
-model_out = f"{model_dir}/model.onnx"
-
-print("## Loading Model")
-tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
-model = AutoModelForSequenceClassification.from_pretrained(model_dir, trust_remote_code=True)
-
-# Set the model in evaluation mode
-model.eval()
-
-# Example input for export
-inputs = tokenizer("Example input text", return_tensors="pt")
-
-
-# Define the export function
-torch.onnx.export(
-    model,                                           # The model to export
-    (inputs["input_ids"], inputs["attention_mask"]), # Model input (you can adjust based on your model's input)
-    model_out,                                       # The path to save the ONNX file
-    export_params=True,                              # Store the trained parameter weights
-    opset_version=14,                                # The ONNX version to use
-    input_names=['input_ids', 'attention_mask'],     # Model's input names
-    output_names=['output'],                         # Model's output names
-    dynamic_axes={'input_ids': {0: 'batch_size'},    # Dynamic axis for input (batch size)
-                  'attention_mask': {0: 'batch_size'},
-                  'output': {0: 'batch_size'}}       # Dynamic axis for output (batch size)
-)
-
-print("## ONNX Model Exported")
-
-# Verify the ONNX model
-print("## Verifying Onnx")
-
-ort_session = ort.InferenceSession(model_out)
-
-if "token_type_ids" in inputs:
-    del inputs["token_type_ids"]
-
-# Prepare inputs for ONNX inference
-ort_inputs = {k: v.cpu().detach().numpy() for k, v in inputs.items()}
-ort_outs = ort_session.run(None, ort_inputs)
-
-print("ONNX output:", ort_outs[0])
-
-with torch.no_grad():
-    pytorch_outputs = model(**inputs)
-    pytorch_output_array = pytorch_outputs.logits.cpu().numpy()
-
-print("PyTorch output:", pytorch_output_array)
-
-# Compare the outputs
-if np.allclose(pytorch_output_array, ort_outs[0], atol=1e-5):
-    print("The ONNX model output matches the PyTorch model output!")
-else:
-    print("The ONNX model output does NOT match the PyTorch model output.")