ai-rec-ner-onnx-model

ONNX-Optimized NER Model for Personal Data Extraction

This repository provides an ONNX-optimized Named Entity Recognition (NER) model for extracting key personal data such as names, cities, and countries from text. The model is optimized for fast inference using ONNX Runtime.

Features

• Optimized for ONNX Runtime.
• Recognizes entities like:
  • First Name
  • Middle Name
  • Last Name
  • City
  • Country
• Easy integration with Hugging Face's optimum library.

Installation

pip install optimum[onnxruntime] transformers

Usage

from typing import Dict, List
from optimum.onnxruntime import ORTModelForTokenClassification
from transformers import AutoTokenizer, pipeline

def extract_entities(text: str) -> Dict[str, str]:
    """
    Extracts personal data entities from a given text using the ONNX model.

    Args:
        text (str): Input text to extract entities from.

    Returns:
        dict: Extracted entities, e.g., {'first_name': 'John', 'city': 'New York'}
    """
    # Load the ONNX model and tokenizer
    model = ORTModelForTokenClassification.from_pretrained(
        "harry-kpv-hf/ai-rec-ner-onnx-model",  # Replace with your Hugging Face model repo
        provider="CPUExecutionProvider"
    )
    tokenizer = AutoTokenizer.from_pretrained("harry-kpv-hf/ai-rec-ner-onnx-model")

    # Create a pipeline for NER
    ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer)

    # Run inference
    ner_results = ner_pipeline(text)

    # Define mapping of entity types to readable keys
    entity_mapping = {
        "U-FIRST_NAME": "first_name",
        "U-MIDDLE_NAME": "middle_name",
        "U-LAST_NAME": "last_name",
        "U-CITY": "city",
        "U-COUNTRY": "country",
    }

    # Process results to create a dictionary of extracted entities
    extracted_data = {}
    for entity in ner_results:
        entity_type = entity["entity"]
        if entity_type in entity_mapping and entity_mapping[entity_type] not in extracted_data:
            extracted_data[entity_mapping[entity_type]] = entity["word"].replace("##", "")

    return extracted_data

# Example usage
text = "John Doe lives in New York, USA."
print(extract_entities(text))
# Output: {'first_name': 'John', 'last_name': 'Doe', 'city': 'New York', 'country': 'USA'}