Vocabulary Pruning Engine for Multilingual GLiNER Models

.gitignore

@@ -27,7 +27,7 @@ eggs/
 .eggs/
 lib/
 lib64/
-parts/
+parts/``
 sdist/
 var/
 wheels/
@@ -92,6 +92,10 @@ target/
 # IPython
 profile_default/
 ipython_config.py
+pruning_adr.md
+CLAUDE.md
+/results/
+ROADMAP.md
 
 # pyenv
 #   For a library or package, you might want to ignore these files since the code is

README.md

@@ -143,6 +143,34 @@ model = GLiNER.from_pretrained(
 
 Find more information on compilation and other optimizations in the [documentation](https://urchade.github.io/GLiNER/usage.html#quantization-compilation-flashdeberta).
 
+## Vocabulary Pruning
+
+Multilingual GLiNER models (mDeBERTa-v3) carry a 250k-token embedding matrix. For a single-language deployment, most of those embeddings are never accessed, creating an unnecessary memory and cold-start bottleneck for edge/CPU deployments. The pruning engine identifies the active token set for a target language, slices `word_embeddings.weight` down to the relevant rows, rebuilds the fast tokenizer, and exports a self-contained model that loads with the standard `GLiNER.from_pretrained()` API — no code changes required in the inference path.
+
+| Metric | Original | Pruned | Δ |
+|---|---|---|---|
+| Vocabulary | 250,105 tokens | 90,840 tokens | **−63.7%** |
+| Model size | 1,155.8 MB | 666.5 MB | **−42.3% (−489 MB)** |
+| Entity F1 | baseline | identical | **0% regression** |
+
+*Benchmarked on `urchade/gliner_multi-v2.1`, English Wikipedia (100k articles).*
+
+```bash
+# Prune to English — one command, no accuracy loss
+python scripts/prune_gliner_vocab.py \
+    --model_id urchade/gliner_multi-v2.1 \
+    --dataset_for_vocab wikipedia \
+    --output_dir ./pruned_en \
+    --lang en
+
+# Verify correctness + measure size reduction
+python scripts/validate_pruned_model.py \
+    --original_model_id urchade/gliner_multi-v2.1 \
+    --pruned_model_dir ./pruned_en
+```
+
+Conservative mode (all seen tokens) is lossless; aggressive mode (`--top_k 30000`) trades ~65% size reduction for minor score shifts near the detection threshold. See [`docs/vocab_pruning.md`](docs/vocab_pruning.md) for the full reference.
+
 ## Serving
 
 For production workloads — high-throughput pipelines, multi-user services, or anywhere you need to go beyond single-process `model.inference()` calls — GLiNER provides a Ray Serve-based serving layer. It adds dynamic batching that automatically groups incoming requests, memory-aware batch sizing that prevents CUDA OOM by calibrating against your GPU, precompiled kernels for common batch sizes to avoid first-call latency, horizontal scaling across multiple GPUs via Ray replicas, and an HTTP API for language-agnostic access.

docs/flash_attention.md

@@ -0,0 +1,133 @@
+# FlashDeBERTa — Fast Attention for GLiNER
+
+## Overview
+
+GLiNER's default DeBERTa v2/v3 backbone uses a custom disentangled relative-position
+attention mechanism that computes a full L×L position-bias matrix. This makes memory and
+compute scale **quadratically** with sequence length, which is the root cause of the
+384-token practical limit and slow inference at longer inputs.
+
+[FlashDeBERTa](https://github.com/Knowledgator/FlashDeBERTa) (Knowledgator, 2024) rewrites
+DeBERTa's disentangled attention with Flash Attention-style IO-aware tiling. The same
+pretrained weights are used — no retraining required.
+
+| Sequence length | Standard DeBERTa | FlashDeBERTa | Speedup |
+|---|---|---|---|
+| 128 tokens | baseline | baseline | ~1.1× |
+| 384 tokens | baseline | — | ~1.5× |
+| 512 tokens | baseline | — | **~2×** |
+| 1024 tokens | baseline | — | **~3–4×** |
+| 2048 tokens | OOM on 8 GB | — | ∞ (enables long context) |
+| 4096 tokens | OOM on 16 GB | — | **~5×** |
+
+> Numbers are approximate and hardware-dependent. Run `scripts/benchmark_flash_attention.py`
+> to get exact figures on your machine.
+
+## Requirements
+
+```bash
+pip install flashdeberta        # requires Python ≥ 3.10
+```
+
+> FlashDeBERTa is only available for **Python 3.10 or later** and applies only to models
+> with a DeBERTa v2/v3 backbone (e.g. `urchade/gliner_multi-v2.1`,
+> `urchade/gliner_large-v2.1`). Models using other backbones (BERT, ModernBERT, T5)
+> are unaffected — `flash_attention=True` is silently ignored for them.
+
+## Usage
+
+### `from_pretrained`
+
+```python
+from gliner import GLiNER
+
+# Enable Flash Attention at load time
+model = GLiNER.from_pretrained(
+    "urchade/gliner_multi-v2.1",
+    flash_attention=True,
+)
+
+entities = model.predict_entities(
+    "Apple Inc. was founded by Steve Jobs in Cupertino.",
+    ["person", "organization", "location"],
+)
+```
+
+### `load_from_config`
+
+```python
+model = GLiNER.load_from_config(
+    "path/to/gliner_config.json",
+    flash_attention=True,
+    backbone_from_pretrained=True,
+)
+```
+
+### Persistent (saved in `gliner_config.json`)
+
+When you call `model.save_pretrained(output_dir)` on a FlashDeBERTa model, the config
+is saved with `"use_flash_attention": true`. The next `from_pretrained(output_dir)` call
+will automatically reload with Flash Attention — no need to pass `flash_attention=True`
+again.
+
+### Environment variable (legacy)
+
+The previous `USE_FLASHDEBERTA=1` environment variable still works as a fallback:
+
+```bash
+USE_FLASHDEBERTA=1 python my_script.py
+```
+
+## Fallback behaviour
+
+If `flash_attention=True` is requested but `flashdeberta` is not installed, GLiNER emits
+a `UserWarning` and falls back to standard DeBERTa attention — the model loads and runs
+correctly, just without the speedup:
+
+```
+UserWarning: use_flash_attention=True requested but 'flashdeberta' is not installed.
+Falling back to standard DeBERTa attention.
+Install with: pip install flashdeberta
+```
+
+## Benchmarking
+
+```bash
+python scripts/benchmark_flash_attention.py \
+    --model_id urchade/gliner_multi-v2.1 \
+    --output_dir results/flash_benchmark \
+    --n_runs 20
+```
+
+Outputs:
+- `results/flash_benchmark/flash_attention_benchmark.csv` — latency table
+- `results/flash_benchmark/flash_attention_benchmark.png` — speedup plot
+
+## Supported architectures
+
+| Architecture | Flash Attention support |
+|---|---|
+| `UniEncoderSpan` (DeBERTa v2/v3 backbone) | ✅ |
+| `UniEncoderToken` (DeBERTa v2/v3 backbone) | ✅ |
+| `BiEncoderSpan` (DeBERTa v2/v3 backbone) | ✅ |
+| Models with BERT, RoBERTa, ModernBERT backbone | ❌ Not applicable — these use standard SDPA |
+| Models with T5/MT5 backbone | ❌ Not applicable |
+
+## Long-context inference
+
+FlashDeBERTa makes sequences beyond 384 tokens practical. Combine with the sliding-window
+utility (`predict_entities_long()`) for full long-document support:
+
+```python
+model = GLiNER.from_pretrained("urchade/gliner_multi-v2.1", flash_attention=True)
+
+# Process a 2000-token document in overlapping 1024-token windows
+entities = model.predict_entities_long(
+    long_document_text,
+    ["person", "organization", "location"],
+    max_tokens=1024,
+    stride=256,
+)
+```
+
+See [`docs/long_document_inference.md`](long_document_inference.md) for the full guide.

docs/index.md

@@ -16,6 +16,9 @@ training
 architectures
 add_custom_architectures
 convert_to_onnx
+vocab_pruning
+flash_attention
+relation_extraction
 serving
 ```
 

docs/relation_extraction.md

@@ -0,0 +1,171 @@
+# Joint NER + Relation Extraction
+
+## Overview
+
+`UniEncoderSpanRelexGLiNER` extends GLiNER's span-based NER architecture with a
+relation extraction head, performing both tasks in a single forward pass. This
+enables knowledge-graph construction, document understanding, and structured
+information extraction without requiring two separate models.
+
+Reference: [GLiNER-Relex (arXiv:2605.10108)](https://arxiv.org/abs/2605.10108)
+
+## Architecture
+
+```
+Text → DeBERTa encoder → Span representations
+                              │
+                              ├── NER head   → entity predictions
+                              │
+                              └── RE head    → entity-pair adjacency
+                                               → relation type classification
+```
+
+The RE head scores all entity-pair combinations using the span representations
+from the NER head. No separate encoding pass is needed.
+
+## Inference
+
+### `predict_entities()` — entities only
+
+```python
+from gliner.model import UniEncoderSpanRelexGLiNER
+
+model = UniEncoderSpanRelexGLiNER.from_pretrained("knowledgator/gliner-relex-large-v1.0")
+
+entities = model.predict_entities(
+    "Apple Inc. was founded by Steve Jobs in Cupertino, California.",
+    labels=["organization", "person", "location"],
+)
+# → [{"text": "Apple Inc.", "label": "organization", ...}, ...]
+```
+
+### `predict_relations()` — entities + relations
+
+```python
+entities, relations = model.predict_relations(
+    "Apple Inc. was founded by Steve Jobs in Cupertino, California.",
+    labels=["organization", "person", "location"],
+    relations=["founded_by", "headquartered_in", "born_in"],
+    threshold=0.5,
+)
+
+for rel in relations:
+    print(f"{rel['head']['text']} --[{rel['relation']}]--> {rel['tail']['text']}  (score={rel['score']:.2f})")
+# → Apple Inc. --[founded_by]--> Steve Jobs  (score=0.87)
+# → Apple Inc. --[headquartered_in]--> Cupertino  (score=0.74)
+```
+
+### API reference
+
+```python
+entities, relations = model.predict_relations(
+    text:                str,
+    labels:              List[str],          # entity type labels
+    relations:           List[str],          # relation type labels
+    threshold:           float = 0.5,        # entity confidence threshold
+    adjacency_threshold: float = None,       # entity-pair adjacency threshold (defaults to threshold)
+    relation_threshold:  float = None,       # relation type threshold (defaults to threshold)
+    flat_ner:            bool = True,
+    multi_label:         bool = False,
+)
+# Returns:
+#   entities: List[Dict] — same format as predict_entities()
+#   relations: List[Dict] — {"head": entity_dict, "relation": str, "tail": entity_dict, "score": float}
+```
+
+> **Description conditioning (Feature 2):** Entity labels and relation types accept full
+> natural-language descriptions via the same API as `predict_entities()`.
+> ```python
+> entities, relations = model.predict_relations(
+>     text,
+>     labels={"organization": "a legally incorporated company or institution", ...},
+>     relations=["founded_by", "headquartered_in"],
+> )
+> ```
+
+### Long-document inference
+
+Use `predict_entities_long()` for the entity extraction step, then pass found entities
+directly to the model for relation scoring:
+
+```python
+from gliner.long_doc import predict_entities_long
+
+# Step 1: extract entities with sliding window
+entities = predict_entities_long(model, long_text, labels, max_tokens=512, stride=128)
+
+# Step 2: run relation scoring on the full entity set (no sliding window needed for RE)
+_, relations = model.predict_relations(long_text, labels, relation_types, threshold=0.5)
+```
+
+## Training
+
+Use `scripts/train_relex.py` to train on any dataset with the joint NER+RE format:
+
+```bash
+python scripts/train_relex.py \
+    --model_id microsoft/deberta-v3-small \
+    --train_data data/conll04_train.json \
+    --eval_data  data/conll04_dev.json \
+    --output_dir results/relex_conll04 \
+    --max_steps 5000 \
+    --batch_size 8
+```
+
+### Training data format
+
+JSON or JSON Lines file. Each example must have `tokenized_text`, `ner`, and `relations` keys:
+
+```json
+{
+  "tokenized_text": ["Apple", "was", "founded", "by", "Steve", "Jobs"],
+  "ner": [
+    [0, 0, "organization"],
+    [4, 5, "person"]
+  ],
+  "relations": [
+    [4, 5, "person", 0, 0, "organization", "founded_by"]
+  ]
+}
+```
+
+Relation tuple format: `[head_start, head_end, head_type, tail_start, tail_end, tail_type, relation_type]`
+
+### Training with hard negatives + contrastive loss
+
+```bash
+python scripts/train_relex.py \
+    --model_id microsoft/deberta-v3-small \
+    --train_data data/conll04_train.json \
+    --output_dir results/relex_conll04_enhanced \
+    --hard_negative_ratio 0.5 \
+    --contrastive_coef 0.1 \
+    --use_curriculum
+```
+
+## Supported datasets
+
+| Dataset | Entity types | Relation types | Size |
+|---|---|---|---|
+| CoNLL04 | 4 (PER, ORG, LOC, OTH) | 5 (Work_For, Kill, OrgBased_In, Live_In, Located_In) | 1,137 train |
+| FewRel | 80 relation types | — | 56,000 train |
+| DocRED | 6 entity types | 96 relation types | 3,053 train |
+| Re-TACRED | 4 entity types | 40 relation types | 68,124 train |
+
+## Recommended pretrained models
+
+| Model | Description |
+|---|---|
+| `knowledgator/gliner-relex-large-v1.0` | Large, high accuracy |
+| `knowledgator/gliner-relex-small-v1.0` | Small, fast |
+
+## Troubleshooting
+
+**No relations predicted:**
+- Lower `relation_threshold` (try 0.3)
+- Ensure both entity types in the relation appear in `labels`
+- Verify the model was trained with relation annotations (not NER-only)
+
+**Entity predictions differ from `predict_entities()`:**
+- This is expected: the RE head sees entity context that can shift span scores slightly
+- Use `return_relations=False` in `inference()` to get NER-only output without RE overhead