Upload to_onnx.py

to_onnx.py

@@ -1,256 +1,64 @@
-import os
 import torch
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
-from onnxruntime.quantization import quantize_dynamic, quantize_static, QuantType
-from onnxruntime.quantization.calibrate import CalibrationDataReader
-import onnx
-import time
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import onnxruntime as ort
 import numpy as np
 
-def ensure_directory(path):
-    """Create directory if it doesn't exist"""
-    abs_path = os.path.abspath(path)
-    if not os.path.exists(abs_path):
-        os.makedirs(abs_path)
-        print(f"Created directory: {abs_path}")
-    return abs_path
-
-def verify_file_exists(file_path, timeout=5):
-    """Verify that a file exists and is not empty"""
-    start_time = time.time()
-    while time.time() - start_time < timeout:
-        if os.path.exists(file_path) and os.path.getsize(file_path) > 0:
-            return True
-        time.sleep(0.1)
-    return False
-
-def export_to_onnx(model, tokenizer, save_path):
-    """Export model to ONNX format"""
+def convert_and_test_onnx(model_name, output_path="language_detection.onnx", test_text="This is a test sentence."):
+    """
+    Converts a Hugging Face model to ONNX, modifies the tokenizer, and tests the ONNX model.
+    """
     try:
-        # Create a dummy input for the model
-        dummy_input = tokenizer("This is a sample input", return_tensors="pt")
-
-        # Export the model to ONNX
-        torch.onnx.export(
-            model,
-            (dummy_input["input_ids"], dummy_input["attention_mask"]),
-            save_path,
-            opset_version=14,
-            input_names=["input_ids", "attention_mask"],
-            output_names=["output"],
-            dynamic_axes={
-                "input_ids": {0: "batch_size"},
-                "attention_mask": {0: "batch_size"},
-                "output": {0: "batch_size"}
-            }
-        )
-
-        # Verify the file was created
-        if verify_file_exists(save_path):
-            print(f"Successfully exported ONNX model to {save_path}")
-            return True
-        else:
-            print(f"Failed to verify ONNX model at {save_path}")
-            return False
-    except Exception as e:
-        print(f"Error exporting to ONNX: {str(e)}")
-        return False
-
-def create_calibration_dataset(tokenizer, max_length=512):
-    """Generate calibration dataset for static quantization with padding"""
-    samples = [
-        "This is an English sentence.",
-        "Dies ist ein deutscher Satz.",
-        "C'est une phrase française.",
-        "Esta es una frase en español.",
-        "这是一个中文句子。",
-        "これは日本語の文章です。"
-    ]
-
-    # Tokenize with padding and truncation
-    encoded_samples = []
-    for text in samples:
-        encoded = tokenizer(
-            text,
-            padding='max_length',
-            max_length=max_length,
-            truncation=True,
-            return_tensors="pt"
-        )
-        encoded_samples.append({
-            'input_ids': encoded['input_ids'],
-            'attention_mask': encoded['attention_mask']
-        })
-
-    return encoded_samples
-
-class CalibrationLoader(CalibrationDataReader):
-    def __init__(self, calibration_data):
-        self.calibration_data = calibration_data
-        self.current_index = 0
-
-    def get_next(self):
-        if self.current_index >= len(self.calibration_data):
-            return None
-
-        current_data = self.calibration_data[self.current_index]
-        self.current_index += 1
-
-        # Ensure we're returning numpy arrays with the correct shape
-        return {
-            'input_ids': current_data['input_ids'].numpy(),
-            'attention_mask': current_data['attention_mask'].numpy()
-        }
-
-    def rewind(self):
-        self.current_index = 0
+        # Load the tokenizer and model
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForSequenceClassification.from_pretrained(model_name)
 
-def export_to_onnx(model, tokenizer, save_path, max_length=512):
-    """Export model to ONNX format with fixed dimensions"""
-    try:
-        # Create a dummy input with fixed dimensions
-        dummy_input = tokenizer(
-            "This is a sample input",
-            padding='max_length',
-            max_length=max_length,
-            truncation=True,
-            return_tensors="pt"
-        )
+        # Modify the tokenizer's normalizer
+        if hasattr(tokenizer.backend_tokenizer.normalizer, "normalizations"):
+            tokenizer.backend_tokenizer.normalizer.normalizations = []
+        tokenizer.save_pretrained("./modified_tokenizer")
 
         # Export the model to ONNX
+        dummy_input = tokenizer("This is a test sentence.", return_tensors="pt")
+
         torch.onnx.export(
             model,
             (dummy_input["input_ids"], dummy_input["attention_mask"]),
-            save_path,
-            opset_version=14,
+            output_path,
             input_names=["input_ids", "attention_mask"],
             output_names=["output"],
             dynamic_axes={
-                "input_ids": {0: "batch_size"},
-                "attention_mask": {0: "batch_size"}
-            }
+                "input_ids": {0: "batch", 1: "sequence"},
+                "attention_mask": {0: "batch", 1: "sequence"},
+                "output": {0: "batch"},
+            },
+            opset_version=14,
         )
 
-        if verify_file_exists(save_path):
-            print(f"Successfully exported ONNX model to {save_path}")
-            return True
-        else:
-            print(f"Failed to verify ONNX model at {save_path}")
-            return False
-    except Exception as e:
-        print(f"Error exporting to ONNX: {str(e)}")
-        return False
+        print(f"Model successfully converted and saved to {output_path}")
 
-def quantize_model(base_onnx_path, onnx_dir, config_name, calibration_dataset=None):
-    """
-    Quantize ONNX model using either dynamic or static quantization.
+        # Test the ONNX model
+        ort_session = ort.InferenceSession(output_path)
+        tokenizer_test = AutoTokenizer.from_pretrained("./modified_tokenizer")
 
-    Args:
-        base_onnx_path (str): Path to the base ONNX model
-        onnx_dir (str): Directory to save quantized models
-        config_name (str): Type of quantization ('dynamic' or 'static')
-        calibration_dataset (list, optional): Dataset for static quantization calibration
-    """
-    try:
-        quantized_model_path = os.path.join(onnx_dir, f"model_{config_name}_quantized.onnx")
-
-        if config_name == "dynamic":
-            print(f"\nPerforming dynamic quantization...")
-            quantize_dynamic(
-                model_input=base_onnx_path,
-                model_output=quantized_model_path,
-                weight_type=QuantType.QUInt8
-            )
+        # Explicitly set return_token_type_ids=False
+        inputs = tokenizer_test(test_text, return_tensors="np", return_token_type_ids=False)
 
-        elif config_name == "static" and calibration_dataset is not None:
-            print(f"\nPerforming static quantization...")
-            calibration_loader = CalibrationLoader(calibration_dataset)
-            quantize_static(
-                model_input=base_onnx_path,
-                model_output=quantized_model_path,
-                calibration_data_reader=calibration_loader,
-                quant_format=QuantType.QUInt8
-            )
+        ort_inputs = {k: v for k, v in inputs.items()}
 
-        else:
-            print(f"Invalid quantization configuration: {config_name}")
-            return False
+        ort_outputs = ort_session.run(None, ort_inputs)
+        logits = ort_outputs[0]
+        predicted_class_id = np.argmax(logits, axis=-1)
 
-        # Verify the quantized model exists
-        if verify_file_exists(quantized_model_path):
-            print(f"Successfully created {config_name} quantized model at {quantized_model_path}")
+        label_list = model.config.id2label
+        predicted_label = label_list[predicted_class_id[0]]
 
-            # Print file sizes for comparison
-            base_size = os.path.getsize(base_onnx_path) / (1024 * 1024)  # Convert to MB
-            quantized_size = os.path.getsize(quantized_model_path) / (1024 * 1024)  # Convert to MB
-
-            print(f"Original model size: {base_size:.2f} MB")
-            print(f"Quantized model size: {quantized_size:.2f} MB")
-            print(f"Size reduction: {((base_size - quantized_size) / base_size * 100):.2f}%")
-
-            return True
-        else:
-            print(f"Failed to verify quantized model at {quantized_model_path}")
-            return False
+        print(f"Test text: {test_text}")
+        print(f"Predicted label: {predicted_label}")
 
     except Exception as e:
-        print(f"Error during {config_name} quantization: {str(e)}")
-        return False
-
-
-def main():
-    # Get absolute paths
-    current_dir = os.path.abspath(os.getcwd())
-    onnx_dir = ensure_directory(os.path.join(current_dir, "onnx"))
-    base_onnx_path = os.path.join(onnx_dir, "model.onnx")
-
-    print(f"Working directory: {current_dir}")
-    print(f"ONNX directory: {onnx_dir}")
-    print(f"Base ONNX model path: {base_onnx_path}")
-
-    # Step 1: Load model and tokenizer
-    print("\nLoading model and tokenizer...")
-    model_name = "alexneakameni/language_detection"
-    model = AutoModelForSequenceClassification.from_pretrained(model_name)
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-    # Get the model's default max_length
-    max_length = tokenizer.model_max_length
-
-    # Step 2: Export base ONNX model
-    if not export_to_onnx(model, tokenizer, base_onnx_path, max_length):
-        print("Failed to export base ONNX model. Exiting.")
-        return
-
-    # Verify the ONNX model
-    try:
-        print(f"Verifying ONNX model at: {base_onnx_path}")
-        onnx_model = onnx.load(base_onnx_path)
-        print("Successfully verified ONNX model")
-    except Exception as e:
-        print(f"Error verifying ONNX model: {str(e)}")
-        return
-
-    # Step 3: Create calibration dataset
-    calibration_dataset = create_calibration_dataset(tokenizer, max_length)
-
-    # Step 4: Create quantized versions
-    print("\nCreating quantized versions...")
-
-    # Dynamic quantization
-    quantize_model(
-        base_onnx_path=base_onnx_path,
-        onnx_dir=onnx_dir,
-        config_name="dynamic"
-    )
-
-    # Static quantization
-    quantize_model(
-        base_onnx_path=base_onnx_path,
-        onnx_dir=onnx_dir,
-        config_name="static",
-        calibration_dataset=calibration_dataset
-    )
+        print(f"Error during conversion or testing: {e}")
 
 if __name__ == "__main__":
-    main()
+    model_name = "dewdev/language_detection"
+    test_text = "मैंने राजा को हिंदी में एक पत्र लिखा।"
+    convert_and_test_onnx(model_name, test_text=test_text)