DRU-RE-EntityPair-TwoHead

DRU-RE-EntityPair-TwoHead is an Arabic relation extraction model built on U4RASD/NeoAraBERT for the WojoodRelations / KnowledgeGraphEval relation inventory.

The model predicts whether a relation exists for a subject/object pair, then predicts the positive relation type with a separate head. It expects the subject span, object span, and coarse entity types to be known.

Architecture

The input sentence is marked with typed subject/object markers and encoded with NeoAraBERT. The pair representation is built from:

1. the contextual CLS vector,
2. the contextual opening subject-marker vector,
3. the contextual opening object-marker vector,
4. the absolute subject/object difference vector,
5. the elementwise subject/object product vector,
6. a 42-dimensional directional one-hot entity-type vector.

These features are concatenated into:

5 * encoder_hidden_size + 42 = 5 * 768 + 42 = 3882

That pair vector is passed through a pair MLP:

3882 -> 1024 -> 512

The resulting shared pair representation feeds two heads:

• a binary existence head: no_relation vs relation_exists
• a 40-way positive-relation head: only positive labels, no no_relation class

Training behavior:

• for no_relation rows, only the existence loss is active
• for positive rows, both the existence loss and positive-label loss are active

Training Setup

• Base encoder: U4RASD/NeoAraBERT
• Dataset: U4RASD/WojoodRelationsAnnotated
• Train file: train.jsonl
• Labeled validation file: val.jsonl
• Official prediction input: val_official_annotated_unlabeled.jsonl
• Max length: 512
• Context chars: 300
• Epochs: 4
• Train batch size: 4
• Eval batch size: 8
• Gradient accumulation: 4
• Effective batch size: 16
• Encoder LR: 1e-5
• Head LR: 5e-5
• Weight decay: 0.01
• Warmup ratio: 0.1
• Pair dropout: 0.1
• Type feature mode: one_hot_42_directional
• Existence loss weight: 1.0
• Positive relation loss weight: 1.0
• Best checkpoint recorded in config:
  • outputs/DRU-RE-EntityPair-TwoHead/checkpoints/checkpoint-1500

Labels And Types

• Total labels: 41
• Positive labels: 40
• Coarse entity types: 21

The config and label/type maps are stored in:

• configs/run_config.json
• configs/architecture_config.json
• configs/labels.json
• configs/positive_labels.json
• configs/entity_types.json
• configs/label2id.json, configs/id2label.json
• configs/positive_label2id.json, configs/positive_id2label.json
• configs/type2id.json, configs/id2type.json

Thresholds

This repo contains two official-style submission variants:

1. tuned-threshold submission:

   • artifacts/submission.zip
   • uses existence threshold 0.14
   • source: auto_tuned_on_labeled_val

2. fixed-threshold submission:

   • artifacts/submission_threshold_0_5.zip
   • uses existence threshold 0.5

The threshold sweep is stored in:

• artifacts/existence_threshold_sweep.csv
• artifacts/best_existence_threshold.json
• artifacts/submission_threshold_0_5_summary.json

From the sweep on labeled val.jsonl:

• best positive micro F1: 0.5429333333333334
• best threshold: 0.14
• labeled-val micro precision: 0.5049603174603174
• labeled-val micro recall: 0.5870818915801614
• labeled-val all-label accuracy: 0.6277415530527564

From the fixed 0.5 summary:

• total rows: 2074
• no_relation predictions: 1322
• changed vs threshold 0.14: 414

Input Contract

Each example is expected to provide:

• sentence
• subject
• object
• subject_start
• subject_end
• object_start
• object_end
• subject_type
• object_type

The official evaluation input should come from:

U4RASD/WojoodRelationsAnnotated/val_official_annotated_unlabeled.jsonl

That file keeps the official validation row set and order, includes spans/types, and intentionally leaves relation blank.

Output Contract

For Codabench-style submission, write:

<triple_id>\t<predicted_relation>

When the internal prediction is no_relation, the submission file should use Codabench's spelling:

no-relation

This repo already includes:

• artifacts/predictions.txt
• artifacts/predictions_threshold_0_5.txt
• artifacts/submission.zip
• artifacts/submission_threshold_0_5.zip
• artifacts/codabench_val_predictions.jsonl
• artifacts/codabench_val_predictions_debug.jsonl

Files

• model/pytorch_model.bin: custom PyTorch state dict
• model/modeling_entity_pair_two_head.py: model definition
• model/config.json: tokenizer/model config metadata
• model/tokenizer.json, model/tokenizer_config.json, model/special_tokens_map.json
• artifacts/labeled_val_metrics.json
• artifacts/labeled_val_classification_report.csv
• artifacts/labeled_val_confusion_matrix.csv
• artifacts/labeled_val_predictions_debug.jsonl
• artifacts/existence_threshold_sweep.csv
• artifacts/submission.zip
• artifacts/submission_threshold_0_5.zip

Evaluation Notes

The repo stores threshold-aware labeled validation metrics in:

• artifacts/labeled_val_metrics.json
• artifacts/existence_threshold_sweep.csv
• artifacts/best_existence_threshold.json

The main model-selection metric recorded in run_config.json is the same value reported in the labeled validation metrics:

best_metric = micro_f1_positive = 0.5429333333333334

This corresponds to the tuned existence threshold 0.14 on labeled val.jsonl.

Practical Use

Use this model when:

• you have reliable subject/object spans,
• you have reliable coarse entity types,
• you want a stronger binary relation detector than the older one-head DRU-RE variants,
• you want access to both a tuned official submission and a stricter 0.5 threshold submission.

Limitations

• This is not a general-purpose Arabic RE model for arbitrary schemas.
• It is tied to the WojoodRelations label inventory and preprocessing assumptions.
• Performance depends heavily on correct spans and correct coarse entity types.
• The existence threshold changes behavior materially. 0.14 and 0.5 are meaningfully different operating points.

Provenance

• Dataset: https://huggingface.co/datasets/U4RASD/WojoodRelationsAnnotated
• Base encoder: https://huggingface.co/U4RASD/NeoAraBERT
• Wojood / ArabicNER: https://github.com/SinaLab/ArabicNER
• WojoodRelations paper: https://aclanthology.org/2025.emnlp-main.1741/
• Sina relation tools: https://sina.birzeit.edu/relations/