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

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

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.

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

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.

Downloads last month
47
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Dataset used to train U4RASD/TypePredictor