Spaces:
Sleeping
Sleeping
| """ | |
| NER Model Inference Wrapper for Route Extraction. | |
| Supports: | |
| - PyTorch model inference | |
| - ONNX Runtime inference (faster, quantized) | |
| - Automatic model loading with fallback | |
| """ | |
| import os | |
| import numpy as np | |
| from typing import List, Tuple, Optional | |
| # ============================================================ | |
| # Label Configuration | |
| # ============================================================ | |
| LABEL_LIST = ['O', 'B-ORIGIN', 'I-ORIGIN', 'B-DEST', 'I-DEST'] | |
| ID2LABEL = {i: label for i, label in enumerate(LABEL_LIST)} | |
| LABEL2ID = {label: i for i, label in enumerate(LABEL_LIST)} | |
| class NERModel: | |
| """ | |
| NER model wrapper supporting PyTorch and ONNX inference. | |
| Usage: | |
| model = NERModel("path/to/model") | |
| entities = model.predict("من مدينة نصر الى المعادي") | |
| # Returns: {'origin': 'مدينة نصر', 'origin_score': 0.95, | |
| # 'destination': 'المعادي', 'destination_score': 0.92} | |
| """ | |
| def __init__(self, model_path: str, use_onnx: bool = False): | |
| """ | |
| Load NER model. | |
| Args: | |
| model_path: Path to saved model directory | |
| use_onnx: If True, load ONNX model (faster inference) | |
| """ | |
| self.model_path = model_path | |
| self.use_onnx = use_onnx | |
| self.model = None | |
| self.tokenizer = None | |
| self.ort_session = None | |
| if use_onnx: | |
| self._load_onnx(model_path) | |
| else: | |
| self._load_pytorch(model_path) | |
| def _load_pytorch(self, model_path: str): | |
| """Load PyTorch model.""" | |
| import torch | |
| from transformers import AutoTokenizer, AutoModelForTokenClassification | |
| print(f"Loading NER model (PyTorch): {model_path}") | |
| self.tokenizer = AutoTokenizer.from_pretrained(model_path) | |
| self.model = AutoModelForTokenClassification.from_pretrained(model_path) | |
| self.model.eval() | |
| self.device = torch.device('cpu') # CPU for production | |
| self.model.to(self.device) | |
| print("NER model loaded successfully (PyTorch)") | |
| def _load_onnx(self, model_path: str): | |
| """Load ONNX model for faster inference.""" | |
| from transformers import AutoTokenizer | |
| print(f"Loading NER model (ONNX): {model_path}") | |
| self.tokenizer = AutoTokenizer.from_pretrained(model_path) | |
| # Try optimum first, fall back to onnxruntime | |
| try: | |
| from optimum.onnxruntime import ORTModelForTokenClassification | |
| self.ort_model = ORTModelForTokenClassification.from_pretrained(model_path) | |
| self._onnx_method = 'optimum' | |
| except ImportError: | |
| import onnxruntime as ort | |
| onnx_path = os.path.join(model_path, 'model.onnx') | |
| self.ort_session = ort.InferenceSession( | |
| onnx_path, | |
| providers=['CPUExecutionProvider'], | |
| ) | |
| self._onnx_method = 'onnxruntime' | |
| print(f"NER model loaded successfully (ONNX via {self._onnx_method})") | |
| def predict(self, text: str) -> dict: | |
| """ | |
| Run NER prediction on text. | |
| Returns dict with: | |
| - origin: extracted origin text (or None) | |
| - origin_score: confidence score (0-1) | |
| - destination: extracted destination text (or None) | |
| - destination_score: confidence score (0-1) | |
| """ | |
| if not text or not text.strip(): | |
| return {'origin': None, 'origin_score': 0, 'destination': None, 'destination_score': 0} | |
| tokens = text.split() | |
| if self.use_onnx: | |
| predicted_labels, scores = self._predict_onnx(tokens) | |
| else: | |
| predicted_labels, scores = self._predict_pytorch(tokens) | |
| # Extract entity spans from BIO labels | |
| origin, origin_score = self._extract_entity(tokens, predicted_labels, scores, 'ORIGIN') | |
| destination, dest_score = self._extract_entity(tokens, predicted_labels, scores, 'DEST') | |
| return { | |
| 'origin': origin, | |
| 'origin_score': origin_score, | |
| 'destination': destination, | |
| 'destination_score': dest_score, | |
| } | |
| def _predict_pytorch(self, tokens: List[str]) -> Tuple[List[str], List[float]]: | |
| """Run PyTorch inference.""" | |
| import torch | |
| encoding = self.tokenizer( | |
| tokens, | |
| is_split_into_words=True, | |
| truncation=True, | |
| max_length=128, | |
| padding=True, | |
| return_tensors='pt', | |
| ) | |
| with torch.no_grad(): | |
| outputs = self.model( | |
| input_ids=encoding['input_ids'].to(self.device), | |
| attention_mask=encoding['attention_mask'].to(self.device), | |
| ) | |
| logits = outputs.logits[0] | |
| probs = torch.softmax(logits, dim=-1) | |
| predictions = torch.argmax(logits, dim=-1) | |
| # Align predictions back to original tokens | |
| word_ids = encoding.word_ids() | |
| labels = [] | |
| scores = [] | |
| prev_word_id = None | |
| for idx, word_id in enumerate(word_ids): | |
| if word_id is not None and word_id != prev_word_id: | |
| pred_id = predictions[idx].item() | |
| labels.append(ID2LABEL[pred_id]) | |
| scores.append(probs[idx][pred_id].item()) | |
| prev_word_id = word_id | |
| return labels[:len(tokens)], scores[:len(tokens)] | |
| def _predict_onnx(self, tokens: List[str]) -> Tuple[List[str], List[float]]: | |
| """Run ONNX inference.""" | |
| encoding = self.tokenizer( | |
| tokens, | |
| is_split_into_words=True, | |
| truncation=True, | |
| max_length=128, | |
| padding=True, | |
| return_tensors='np' if self._onnx_method == 'onnxruntime' else 'pt', | |
| ) | |
| if self._onnx_method == 'optimum': | |
| inputs = { | |
| 'input_ids': encoding['input_ids'], | |
| 'attention_mask': encoding['attention_mask'], | |
| } | |
| if 'token_type_ids' in encoding: | |
| inputs['token_type_ids'] = encoding['token_type_ids'] | |
| outputs = self.ort_model(**inputs) | |
| logits = outputs.logits[0].detach().numpy() | |
| else: | |
| ort_inputs = { | |
| 'input_ids': encoding['input_ids'].astype(np.int64), | |
| 'attention_mask': encoding['attention_mask'].astype(np.int64), | |
| } | |
| if 'token_type_ids' in encoding: | |
| ort_inputs['token_type_ids'] = encoding['token_type_ids'].astype(np.int64) | |
| logits = self.ort_session.run(None, ort_inputs)[0][0] | |
| # Softmax | |
| exp_logits = np.exp(logits - np.max(logits, axis=-1, keepdims=True)) | |
| probs = exp_logits / np.sum(exp_logits, axis=-1, keepdims=True) | |
| predictions = np.argmax(logits, axis=-1) | |
| # Align predictions back to original tokens | |
| word_ids = encoding.word_ids() | |
| labels = [] | |
| scores = [] | |
| prev_word_id = None | |
| for idx, word_id in enumerate(word_ids): | |
| if word_id is not None and word_id != prev_word_id: | |
| pred_id = int(predictions[idx]) | |
| labels.append(ID2LABEL[pred_id]) | |
| scores.append(float(probs[idx][pred_id])) | |
| prev_word_id = word_id | |
| return labels[:len(tokens)], scores[:len(tokens)] | |
| def _extract_entity(self, tokens: List[str], labels: List[str], | |
| scores: List[float], entity_type: str) -> Tuple[Optional[str], float]: | |
| """Extract entity span from BIO-tagged tokens.""" | |
| entity_tokens = [] | |
| entity_scores = [] | |
| in_entity = False | |
| for token, label, score in zip(tokens, labels, scores): | |
| if label == f'B-{entity_type}': | |
| entity_tokens = [token] | |
| entity_scores = [score] | |
| in_entity = True | |
| elif label == f'I-{entity_type}' and in_entity: | |
| entity_tokens.append(token) | |
| entity_scores.append(score) | |
| else: | |
| if in_entity: | |
| break # End of entity | |
| in_entity = False | |
| if entity_tokens: | |
| entity_text = ' '.join(entity_tokens) | |
| avg_score = sum(entity_scores) / len(entity_scores) | |
| return entity_text, avg_score | |
| return None, 0.0 | |