---
language:
  - ar
tags:
  - arabic
  - relation-extraction
  - wojood
  - neoarabert
library_name: pytorch
base_model: U4RASD/NeoAraBERT
datasets:
  - U4RASD/WojoodRelationsAnnotated
metrics:
  - f1
  - precision
  - recall
---

# 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:

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

That pair vector is passed through a pair MLP:

```text
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:

```text
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:

```text
<triple_id>\t<predicted_relation>
```

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

```text
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:

```text
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/