README.md

---
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`.