Spaces:
Sleeping
Sleeping
File size: 1,368 Bytes
9548e93 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 | """
ClauseGuard — Export fine-tuned Legal-BERT to ONNX for fast inference.
Requires: pip install optimum[onnxruntime]
"""
import os
import sys
MODEL_PATH = os.environ.get("MODEL_PATH", "./clauseguard-model/final")
ONNX_OUTPUT = os.environ.get("ONNX_OUTPUT", "./clauseguard-model-onnx")
print(f"📦 Exporting {MODEL_PATH} → ONNX at {ONNX_OUTPUT}")
try:
from optimum.onnxruntime import ORTModelForSequenceClassification
from transformers import AutoTokenizer
# Load PyTorch model and export to ONNX
model = ORTModelForSequenceClassification.from_pretrained(MODEL_PATH, export=True)
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
# Save ONNX model + tokenizer + config
model.save_pretrained(ONNX_OUTPUT)
tokenizer.save_pretrained(ONNX_OUTPUT)
print(f"✅ ONNX model saved to {ONNX_OUTPUT}")
print(f" Files: {os.listdir(ONNX_OUTPUT)}")
# Verify inference works
from transformers import pipeline
classifier = pipeline("text-classification", model=model, tokenizer=tokenizer, top_k=None)
test = classifier("The company may terminate your account at any time without notice.")
print(f" Test inference: {test}")
except ImportError:
print("❌ Install optimum: pip install optimum[onnxruntime]")
sys.exit(1)
except Exception as e:
print(f"❌ Export failed: {e}")
sys.exit(1)
|