cf8f5a7232
Introduce dual-script ONNX NER routing (English/Tamil/mixed), Tamil console samples and integration tests, model download scripts, and a resizable WPF MVVM harness with click-to-load prompts, batch validation, and runtime-adjustable detection panels.
149 lines
4.6 KiB
Python
149 lines
4.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Download and export prachuryyaIITG/SampurNER_Tamil_IndicBERTv2 to ONNX for the PII Redaction POC."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parent.parent
|
|
MODELS_DIR = REPO_ROOT / "models" / "ta"
|
|
MODEL_ID = "prachuryyaIITG/SampurNER_Tamil_IndicBERTv2"
|
|
REQUIRED_PACKAGES = ("transformers", "optimum[onnxruntime]", "onnx", "torch")
|
|
SENTENCEPIECE_CANDIDATES = (
|
|
"sentencepiece.bpe.model",
|
|
"spiece.model",
|
|
"tokenizer.model",
|
|
)
|
|
TOKENIZER_JSON = "tokenizer.json"
|
|
VOCAB_TXT = "vocab.txt"
|
|
|
|
|
|
def ensure_dependencies() -> None:
|
|
try:
|
|
import optimum.onnxruntime # noqa: F401
|
|
import transformers # noqa: F401
|
|
except ImportError:
|
|
print("Installing Python dependencies (this may take a few minutes)...")
|
|
subprocess.check_call(
|
|
[sys.executable, "-m", "pip", "install", *REQUIRED_PACKAGES],
|
|
stdout=sys.stdout,
|
|
stderr=sys.stderr,
|
|
)
|
|
|
|
|
|
def copy_sentencepiece_model(source_dir: Path, target_dir: Path) -> Path | None:
|
|
for name in SENTENCEPIECE_CANDIDATES:
|
|
candidate = source_dir / name
|
|
if candidate.exists():
|
|
destination = target_dir / name
|
|
shutil.copy(candidate, destination)
|
|
return destination
|
|
|
|
return None
|
|
|
|
|
|
def export_wordpiece_assets(source_dir: Path, target_dir: Path) -> list[Path]:
|
|
saved: list[Path] = []
|
|
tokenizer_json = source_dir / TOKENIZER_JSON
|
|
if not tokenizer_json.exists():
|
|
return saved
|
|
|
|
destination = target_dir / TOKENIZER_JSON
|
|
shutil.copy(tokenizer_json, destination)
|
|
saved.append(destination)
|
|
|
|
with tokenizer_json.open(encoding="utf-8") as tokenizer_file:
|
|
tokenizer_data = json.load(tokenizer_file)
|
|
|
|
vocab = tokenizer_data.get("model", {}).get("vocab")
|
|
if not isinstance(vocab, dict):
|
|
return saved
|
|
|
|
vocab_path = target_dir / VOCAB_TXT
|
|
ordered_tokens = [token for token, _ in sorted(vocab.items(), key=lambda item: item[1])]
|
|
vocab_path.write_text("\n".join(ordered_tokens), encoding="utf-8")
|
|
saved.append(vocab_path)
|
|
return saved
|
|
|
|
|
|
def copy_tokenizer_assets(source_dir: Path, target_dir: Path) -> list[Path]:
|
|
sentencepiece_path = copy_sentencepiece_model(source_dir, target_dir)
|
|
if sentencepiece_path is not None:
|
|
return [sentencepiece_path]
|
|
|
|
wordpiece_assets = export_wordpiece_assets(source_dir, target_dir)
|
|
if wordpiece_assets:
|
|
print(
|
|
"WARNING: Hugging Face repo has no SentencePiece model; saved WordPiece "
|
|
f"assets ({', '.join(path.name for path in wordpiece_assets)}). "
|
|
"TamilOnnxNerRunner uses vocab.txt (WordPiece) when present, "
|
|
"otherwise SentencePiece model files."
|
|
)
|
|
return wordpiece_assets
|
|
|
|
raise FileNotFoundError(
|
|
f"No tokenizer assets found in {source_dir}. "
|
|
f"Expected one of {SENTENCEPIECE_CANDIDATES} or {TOKENIZER_JSON}."
|
|
)
|
|
|
|
|
|
def export_model() -> None:
|
|
from optimum.onnxruntime import ORTModelForTokenClassification
|
|
from transformers import AutoTokenizer
|
|
|
|
MODELS_DIR.mkdir(parents=True, exist_ok=True)
|
|
temp_dir = MODELS_DIR / "_export_temp"
|
|
if temp_dir.exists():
|
|
shutil.rmtree(temp_dir)
|
|
temp_dir.mkdir()
|
|
|
|
print(f"Exporting {MODEL_ID} to ONNX...")
|
|
model = ORTModelForTokenClassification.from_pretrained(MODEL_ID, export=True)
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
|
|
|
|
model.save_pretrained(temp_dir)
|
|
tokenizer.save_pretrained(temp_dir)
|
|
|
|
onnx_files = sorted(temp_dir.glob("*.onnx"))
|
|
if not onnx_files:
|
|
raise FileNotFoundError("Export completed but no .onnx file was produced.")
|
|
|
|
target_onnx = MODELS_DIR / "model.onnx"
|
|
shutil.copy(onnx_files[0], target_onnx)
|
|
|
|
config_path = temp_dir / "config.json"
|
|
with config_path.open(encoding="utf-8") as config_file:
|
|
config = json.load(config_file)
|
|
|
|
id2label = config.get("id2label", {})
|
|
labels = [id2label[str(index)] for index in range(len(id2label))]
|
|
labels_path = MODELS_DIR / "ner-labels.txt"
|
|
labels_path.write_text("\n".join(labels), encoding="utf-8")
|
|
|
|
tokenizer_assets = copy_tokenizer_assets(temp_dir, MODELS_DIR)
|
|
shutil.rmtree(temp_dir)
|
|
|
|
print()
|
|
print("Tamil NER model assets saved:")
|
|
print(f" {target_onnx}")
|
|
for asset in tokenizer_assets:
|
|
print(f" {asset}")
|
|
print(f" {labels_path}")
|
|
print()
|
|
print("Run from repository root:")
|
|
print(" dotnet test --filter Category=TamilNer")
|
|
|
|
|
|
def main() -> int:
|
|
ensure_dependencies()
|
|
export_model()
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|