--- language: - ar library_name: pytorch pipeline_tag: text-classification tags: - arabic - named-entity-recognition - entity-typing - wojood - neoarabert - mention-classification datasets: - U4RASD/TypePrediction metrics: - accuracy - f1 --- # TypePredictor `TypePredictor` is a mention-level Arabic entity type classifier. It assumes the entity span is already known, inserts `[ENT]` and `[/ENT]` around that span, and predicts one of 21 Wojood-style entity types. This model is intended as a type-normalization component for a later relation-extraction pipeline, where relation subjects and objects are already available as spans and need a consistent coarse entity type. ## Key result The repository root contains the best checkpoint selected by overall validation macro F1. The best checkpoint was step `25,000` / epoch `3.968292`. | Split | Rows | Accuracy | Micro F1 | Macro F1 present types | Macro F1 all 21 | Weighted F1 | |---|---:|---:|---:|---:|---:|---:| | Validation overall | 12,600 | 0.977857 | 0.977857 | 0.958728 | 0.958728 | 0.977791 | | Test overall | 12,600 | 0.979365 | 0.979365 | 0.959920 | 0.959920 | 0.979297 | ## What this model does and does not do - It classifies a supplied mention/span. - It does not detect entity boundaries. - It does not include a `NONE` class because the training, validation, and test splits do not contain `NONE` examples. - It uses one shared classifier for all mentions; there is no separate subject or object head. ## Architecture ```text Arabic sentence + known character span -> insert [ENT] and [/ENT] around the exact span -> tokenize with U4RASD/NeoAraBERT tokenizer -> U4RASD/NeoAraBERT encoder -> final hidden state at CLS position -> Dropout(0.10) -> Linear(768, 21) -> argmax entity type ``` Architecture details: - Base encoder: `U4RASD/NeoAraBERT` - Hidden size: `768` - Tokenizer size after markers: `65,002` - Classifier shape: `768 -> 21` - Total parameters: `248,162,325` - Encoder parameters: `248,146,176` - Classifier parameters: `16,149` - Loss: ordinary unweighted multiclass cross-entropy - No class weights, focal loss, oversampling, weighted sampler, span pooling, threshold, extra MLP, or two-head design. ## Labels GPE, ORG, DATE, PERS, NORP, ORDINAL, OCC, EVENT, CARDINAL, LOC, WEBSITE, FAC, LAW, TIME, MONEY, CURR, LANGUAGE, PERCENT, PRODUCT, QUANTITY, UNIT ## Dataset - Dataset repository: `U4RASD/TypePrediction` - Train file: `type_predictor_train.jsonl` - Validation file: `type_predictor_val.jsonl` - Test file: `type_predictor_test.jsonl` - Train rows: `100,796` - Validation rows: `12,600` - Test rows: `12,600` - Dataset validation timestamp: `2026-07-12T15:30:12+00:00` - IDs unique across splits: `True` The split is mention-level. A sentence can appear in more than one split with a different target mention. For that reason, the model card reports both: - `seen_sentence_new_entity`: the sentence text was seen in training, but the evaluated target mention is new. - `unseen_sentence`: the sentence text was not seen in training. ## Preprocessing and encoding - Markers: `[ENT]` and `[/ENT]` - Marker insertion is based on exact character offsets, not string replacement. - Maximum sequence length: `512` - Default context window: `300` characters - Fallback context candidates: `None, 500, 300, 150, 80, 30, 0` - The encoder rejects examples where truncation fails to preserve exactly one opening marker and one closing marker in the correct order. ## Training configuration | Setting | Value | |---|---:| | Seed | `42` | | Epochs | `4.0` | | Train batch size/device | `4` | | Eval batch size/device | `8` | | Gradient accumulation | `4` | | Effective batch size | `16` | | Encoder learning rate | `1e-05` | | Classifier learning rate | `5e-05` | | Weight decay | `0.01` | | Warmup ratio | `0.1` | | Max grad norm | `1.0` | | Dropout | `0.1` | | Logging steps | `50` | | Eval steps | `500` | | Save steps | `500` | | FP16 | `True` | | Best model criterion | overall validation macro F1 | Training runtime: - Started: `2026-07-12T15:30:38+00:00` - Finished: `2026-07-12T18:55:01+00:00` - Wall time: `12262.47` seconds (`3.41` hours) - Train runtime reported by Trainer: `12261.72` seconds - Train samples/sec: `32.882` - Train steps/sec: `2.055` - Final train loss: `0.190693` Hardware/runtime: - GPU: `NVIDIA A40` - GPU VRAM: `44.43 GiB` - PyTorch: `2.5.1+cu124` - Transformers: `4.49.0` - Python: `3.12.3` ## Validation curve Validation ran every `500` optimizer steps. The curve below was generated from the Trainer eval log and is included in this repository as `figures/eval_metrics.svg`; the raw curve data is included as `figures/eval_metrics.csv` and `figures/eval_metrics.json`. ![Validation metrics over training](figures/eval_metrics.svg) ## Checkpoints - Best checkpoint source: `outputs/TypePredictor/checkpoints/checkpoint-25000` - Preserved best checkpoint: `outputs/TypePredictor/best_checkpoint` - Best step: `25,000` - Best epoch: `3.968292` - Best validation macro F1: `0.958728` - Latest completed checkpoint source: `outputs/TypePredictor/checkpoints/checkpoint-25196` - Preserved latest checkpoint: `outputs/TypePredictor/latest_checkpoint` - Latest step: `25,196` - Latest epoch: `3.999405` - Repository root released model: `best checkpoint` - `checkpoints/latest/` contains the latest completed checkpoint snapshot. ## Validation results | Subset | Rows | Accuracy | Micro F1 | Macro F1 present types | Macro F1 all 21 | Weighted F1 | |---|---:|---:|---:|---:|---:|---:| | overall | 12,600 | 0.977857 | 0.977857 | 0.958728 | 0.958728 | 0.977791 | | unseen_sentence | 939 | 0.945687 | 0.945687 | 0.908841 | 0.865563 | 0.945011 | | seen_sentence_new_entity | 11,661 | 0.980448 | 0.980448 | 0.963262 | 0.963262 | 0.980403 | ## Test results | Subset | Rows | Accuracy | Micro F1 | Macro F1 present types | Macro F1 all 21 | Weighted F1 | |---|---:|---:|---:|---:|---:|---:| | overall | 12,600 | 0.979365 | 0.979365 | 0.959920 | 0.959920 | 0.979297 | | unseen_sentence | 994 | 0.962777 | 0.962777 | 0.893174 | 0.850642 | 0.961780 | | seen_sentence_new_entity | 11,606 | 0.980786 | 0.980786 | 0.964113 | 0.964113 | 0.980712 | ## Per-class validation results ### Validation overall | Label | Precision | Recall | F1 | Support | |---|---:|---:|---:|---:| | GPE | 0.977432 | 0.982156 | 0.979789 | 2,690 | | ORG | 0.972832 | 0.982707 | 0.977745 | 2,660 | | DATE | 0.993431 | 0.994939 | 0.994185 | 1,976 | | PERS | 0.973832 | 0.972015 | 0.972923 | 1,072 | | NORP | 0.975050 | 0.955969 | 0.965415 | 1,022 | | ORDINAL | 0.979513 | 0.980769 | 0.980141 | 780 | | OCC | 0.981889 | 0.986996 | 0.984436 | 769 | | EVENT | 0.967302 | 0.941645 | 0.954301 | 377 | | CARDINAL | 0.973333 | 0.978552 | 0.975936 | 373 | | LOC | 0.974576 | 0.962343 | 0.968421 | 239 | | WEBSITE | 0.986577 | 0.993243 | 0.989899 | 148 | | FAC | 0.916667 | 0.930769 | 0.923664 | 130 | | LAW | 1.000000 | 1.000000 | 1.000000 | 90 | | TIME | 0.974359 | 0.873563 | 0.921212 | 87 | | MONEY | 0.933333 | 1.000000 | 0.965517 | 42 | | CURR | 1.000000 | 0.951220 | 0.975000 | 41 | | LANGUAGE | 0.933333 | 0.848485 | 0.888889 | 33 | | PERCENT | 1.000000 | 0.903226 | 0.949153 | 31 | | PRODUCT | 1.000000 | 0.842105 | 0.914286 | 19 | | QUANTITY | 1.000000 | 0.818182 | 0.900000 | 11 | | UNIT | 0.909091 | 1.000000 | 0.952381 | 10 | ## Per-class test results ### Test overall | Label | Precision | Recall | F1 | Support | |---|---:|---:|---:|---:| | GPE | 0.973897 | 0.984392 | 0.979117 | 2,691 | | ORG | 0.980769 | 0.978187 | 0.979477 | 2,659 | | DATE | 0.990438 | 0.995448 | 0.992936 | 1,977 | | PERS | 0.979535 | 0.983193 | 0.981361 | 1,071 | | NORP | 0.976540 | 0.978452 | 0.977495 | 1,021 | | ORDINAL | 0.982097 | 0.983355 | 0.982726 | 781 | | OCC | 0.987047 | 0.989610 | 0.988327 | 770 | | EVENT | 0.986413 | 0.962865 | 0.974497 | 377 | | CARDINAL | 0.977901 | 0.951613 | 0.964578 | 372 | | LOC | 0.969565 | 0.933054 | 0.950959 | 239 | | WEBSITE | 0.966887 | 0.986486 | 0.976589 | 148 | | FAC | 0.928000 | 0.899225 | 0.913386 | 129 | | LAW | 1.000000 | 1.000000 | 1.000000 | 91 | | TIME | 0.915663 | 0.873563 | 0.894118 | 87 | | MONEY | 0.973684 | 0.880952 | 0.925000 | 42 | | CURR | 0.975610 | 0.975610 | 0.975610 | 41 | | LANGUAGE | 0.875000 | 0.848485 | 0.861538 | 33 | | PERCENT | 0.964286 | 0.870968 | 0.915254 | 31 | | PRODUCT | 1.000000 | 0.947368 | 0.972973 | 19 | | QUANTITY | 0.909091 | 1.000000 | 0.952381 | 10 | | UNIT | 1.000000 | 1.000000 | 1.000000 | 11 | ### Test unseen-sentence per-class results This subset is the strictest split because the full sentence is unseen during training. Some rare labels have very small support here, so their F1 scores are high variance. ### Test unseen_sentence | Label | Precision | Recall | F1 | Support | |---|---:|---:|---:|---:| | GPE | 0.953333 | 0.934641 | 0.943894 | 153 | | ORG | 0.964029 | 0.964029 | 0.964029 | 139 | | DATE | 0.959677 | 0.991667 | 0.975410 | 120 | | PERS | 0.978873 | 0.972028 | 0.975439 | 143 | | NORP | 0.960317 | 0.968000 | 0.964143 | 125 | | ORDINAL | 0.976190 | 1.000000 | 0.987952 | 82 | | OCC | 1.000000 | 0.960000 | 0.979592 | 50 | | EVENT | 1.000000 | 0.928571 | 0.962963 | 14 | | CARDINAL | 0.979167 | 0.959184 | 0.969072 | 49 | | LOC | 0.933333 | 0.933333 | 0.933333 | 15 | | WEBSITE | 0.945455 | 1.000000 | 0.971963 | 52 | | FAC | 0.888889 | 1.000000 | 0.941176 | 8 | | LAW | 1.000000 | 1.000000 | 1.000000 | 8 | | TIME | 0.900000 | 0.947368 | 0.923077 | 19 | | MONEY | 0.000000 | 0.000000 | 0.000000 | 1 | | CURR | 1.000000 | 0.666667 | 0.800000 | 3 | | LANGUAGE | 0.666667 | 0.500000 | 0.571429 | 8 | | PERCENT | 1.000000 | 1.000000 | 1.000000 | 2 | | PRODUCT | 1.000000 | 1.000000 | 1.000000 | 2 | | QUANTITY | 1.000000 | 1.000000 | 1.000000 | 1 | | UNIT | 0.000000 | 0.000000 | 0.000000 | 0 | ## Confusion-matrix observations The strongest remaining confusions are mostly between semantically adjacent coarse types or rare labels with limited support. ### Top test-overall confusions | Gold label | Predicted label | Count | |---|---|---:| | ORG | GPE | 36 | | GPE | ORG | 21 | | LOC | GPE | 11 | | FAC | GPE | 9 | | EVENT | ORG | 8 | | CARDINAL | ORDINAL | 8 | | NORP | ORG | 7 | | GPE | NORP | 6 | | ORG | NORP | 6 | | ORG | PERS | 5 | | NORP | PERS | 5 | | ORDINAL | ORG | 5 | ### Top test-unseen-sentence confusions | Gold label | Predicted label | Count | |---|---|---:| | GPE | ORG | 3 | | ORG | GPE | 3 | | GPE | WEBSITE | 2 | | PERS | GPE | 2 | | CARDINAL | ORDINAL | 2 | | LANGUAGE | NORP | 2 | | GPE | DATE | 1 | | GPE | PERS | 1 | | GPE | NORP | 1 | | GPE | CARDINAL | 1 | | GPE | LOC | 1 | | ORG | WEBSITE | 1 | ## Artifacts included in this repository - `pytorch_model.bin`: best checkpoint model weights - `config.json`, `tokenizer.json`, `tokenizer_config.json`, `special_tokens_map.json` - `type_predictor_config.json`: task-specific architecture and label config - `modeling_type_predictor.py`: custom PyTorch model wrapper - `inference.py`: local inference example - `metrics/`: validation/test overall, category-specific, and per-class metrics - `confusion_matrices/`: validation/test confusion matrices - `predictions/`: row-level validation and test predictions - `configs/`: architecture, labels, run config, and training arguments - `checkpoint_summary.json`, `run_summary.json`, `evaluation_results.json` - `checkpoints/latest/`: latest completed checkpoint snapshot ## Loading Because this is a small custom wrapper around NeoAraBERT, load the model through the included `modeling_type_predictor.py`. ```python from huggingface_hub import hf_hub_download import importlib.util source = hf_hub_download("U4RASD/TypePredictor", "modeling_type_predictor.py") spec = importlib.util.spec_from_file_location("modeling_type_predictor", source) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) model, tokenizer, config = module.NeoAraBERTTypePredictor.from_pretrained( "U4RASD/TypePredictor" ) model.eval() ``` ## Inference example ```python import torch sentence = "زار أحمد القاهرة أمس." entity = "القاهرة" start = sentence.index(entity) end = start + len(entity) marked = sentence[:start] + " [ENT] " + sentence[start:end] + " [/ENT] " + sentence[end:] batch = tokenizer( marked, return_tensors="pt", truncation=True, max_length=config["max_length"], ) with torch.no_grad(): logits = model(**batch)["logits"] probs = torch.softmax(logits, dim=-1)[0] label_id = int(probs.argmax()) print(config["labels"][label_id], float(probs[label_id])) ``` The repository also includes `inference.py`, which handles marker-preserving entity-centered truncation with the same context-candidate policy used during training. ## Limitations - This is not a full NER model; it requires a known span. - There is no `NONE`/non-entity class in this training run. - Overall metrics are not a pure unseen-sentence estimate because the split is mention-level and intentionally contains sentence overlap. Use `unseen_sentence` metrics for the stricter generalization view. - Rare labels such as `UNIT`, `QUANTITY`, `PRODUCT`, `PERCENT`, and `LANGUAGE` have much lower support than GPE/ORG/DATE/PERS/NORP. - The first locked experiment intentionally avoided imbalance correction, extra classifier layers, span pooling, or threshold tuning. ## Reproducibility notes The generated artifacts contain the run configuration, training arguments, Trainer log history, dataset validation report, encoding validation report, metrics, predictions, and confusion matrices. The test set was evaluated only after training and checkpoint selection; it was not used for checkpoint or hyperparameter selection. Final model card generated from local RunPod artifacts at `2026-07-12T19:07:42+00:00`.