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.

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

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.

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

Test results

Per-class validation results

Validation overall

Per-class test results

Test overall

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

Confusion-matrix observations

The strongest remaining confusions are mostly between semantically adjacent coarse types or rare labels with limited support.

Top test-overall confusions

Top test-unseen-sentence confusions

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.