Repository: 8400TheHealthNetwork/HebSpacy (language: python)

Issue:
run the model with gpu
if i doing prefer_gpu() before the nlp vraible i got error the he need get numpy array in the head_ner in line 42.
How do I overcome this problem?


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/python-app.yml b/.github/workflows/python-app.yml
new file mode 100644
index 0000000..be5178a
--- /dev/null
+++ b/.github/workflows/python-app.yml
@@ -0,0 +1,36 @@
+# This workflow will install Python dependencies, run tests and lint with a single version of Python
+# For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions
+
+name: Python application
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Python 3.8
+      uses: actions/setup-python@v2
+      with:
+        python-version: "3.8"
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install flake8 pytest
+        if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
+    - name: Lint with flake8
+      run: |
+        # stop the build if there are Python syntax errors or undefined names
+        flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
+        # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
+        flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
+    - name: Test with pytest
+      run: |
+        pytest
diff --git a/README.md b/README.md
index 86ebcfe..3b64d44 100644
--- a/README.md
+++ b/README.md
@@ -1,2 +1,80 @@
-# HebSpacy
-Hebrew oriented NER spaCy pipeline
+# HebSpaCy 
+
+A custom spaCy pipeline for Hebrew text including a transformer-based multitask NER model that recognizes 16 entity types in Hebrew, including `GPE`, `PER`, `LOC` and `ORG`.
+
+----
+[![MIT license](https://img.shields.io/badge/license-MIT-brightgreen.svg)](http://opensource.org/licenses/MIT) ![Release](https://img.shields.io/github/v/release/8400TheHealthNetwork/HebSpacy.svg) [![PyPI version](https://badge.fury.io/py/hebspacy.svg)](https://badge.fury.io/py/hebspacy) [![Pypi Downloads](https://img.shields.io/pypi/dm/hebspacy.svg)](https://img.shields.io/pypi/dm/hebspacy.svg) 
+
+## Installation
+
+To run the package you will need to install the package as well as the model, preferably in a virtual environment:
+
+``` sh
+# Create conda env (optional)
+conda create --name hebspacy python=3.8
+conda activate hebsafeharbor
+
+# Install hebspacy
+pip install hebspacy
+
+# Download and install the model (see below availbable models)
+pip install </path/to/download>
+```
+
+#### Available Models
+| Model | Description | Install URL |
+| ----- | ----------- | ----------- |
+| he_ner_news_trf | A full spaCy pipeline for Hebrew text including a multitask NER model trained against the BMC and NEMO corpora. Read more [here](#he_ner_news_trf).| [Download](https://github.com/8400TheHealthNetwork/HebSpacy/releases/download/he_ner_news_trf-3.2.1/he_ner_news_trf-3.2.1-py3-none-any.whl)
+
+
+## Getting started
+```python
+import spacy
+
+nlp = spacy.load("he_ner_news_trf")
+text = """מרגלית דהן
+מספר זהות 11278904-5
+
+2/12/2001
+ביקור חוזר מ18.11.2001
+במסגרת בירור פלפיטציות ואי סבילות למאמצים,מנורגיות קשות ע"ר שרירנים- ביצעה מעבדה שהדגימה:
+המוגלובין 9, מיקרוציטי היפוכרומטי עם RDW 19,
+פריטין 10, סטורציית טרנספרין 8%. 
+מבחינת עומס נגיפי HIV- undetectable ומקפידה על HAART
+"""
+
+doc = nlp(text)
+for entity in doc.ents:
+    print(f"{entity.text} \t {entity.label_}: {entity._.confidence_score:.4f} ({entity.start_char},{entity.end_char})")
+
+>>> מרגלית דהן	 PERS: 0.9999 (0,10)
+>>> 2/12/2001 	 DATE: 0.9897 (33,42)
+>>> מ18.11.2001 	 DATE: 0.8282 (54,65)
+>>> 8% 	 PERCENT: 0.9932 (230,232)
+```
+
+---------------
+### he_ner_news_trf
+'he_ner_news_trf' is a multitask model constructed from [AlephBert](https://arxiv.org/pdf/2104.04052.pdf) and two NER focused heads, each trained against a different NER-annotated Hebrew corpus:
+1. [NEMO corpus](https://github.com/OnlpLab/NEMO-Corpus) - annotations of the Hebrew Treebank (Haaretz newspaper) for the widely-used OntoNotes entity category: `GPE` (geo-political entity), `PER` (person), `LOC` (location), `ORG` (organization), `FAC` (facility), `EVE` (event), `WOA` (work-of-art), `ANG` (language), `DUC` (product). 
+2. [BMC corpus](https://www.cs.bgu.ac.il/~elhadad/nlpproj/naama/) - annotations of articles from Israeli newspapers and websites (Haaretz newspaper, Maariv newspaper, Channel 7) for the common entity categories: `PERS` (person), `LOC` (location), `ORG` (organization), `DATE` (date), `TIME` (time), `MONEY` (money), `PERCENT` (percent), `MISC__AFF` (misc affiliation), `MISC__ENT` (misc entity),
+ `MISC_EVENT` (misc event).
+
+The model was developed and trained using the Hugging Face and PyTorch libraries, and was later integrated into a spaCy pipeline. 
+
+#### Model integration
+The output model was split into three weight files: _the transformer embeddings, the BMC head, and the NEMO head_.
+The components were each packaged in a separate pipe and integrated into the custom pipeline. 
+Furthermore, a custom NER head consolidation pipe was added last to address signal conflicts/overlaps, and sets the `Doc.ents` property.
+
+To access the entities recognized by each NER head, use the `Doc._.<ner_head>` property (e.g., `doc._.nemo_ents` and `doc._.bmc_ents`).
+
+--------------
+## Contribution
+You are welcome to contribute to `hebspacy` project and introduce new feature/ models. 
+Kindly follow the [pipeline codebase instructions](contribute/pipeline/README.md) and the [model training and packaging guidelines](contribute/model/README.md).
+
+
+-----
+
+HebSpaCy is an open-source project developed by [8400 The Health Network](https://www.8400thn.org/).
diff --git a/contribute/model/README.md b/contribute/model/README.md
new file mode 100644
index 0000000..38ae328
--- /dev/null
+++ b/contribute/model/README.md
@@ -0,0 +1,70 @@
+# Contribute your own `hebspacy` model
+You are welcome to contribute by training and packaging your own `hebspacy` model. Please follow the instructions below to enable seamless loading process.
+
+## Model Training
+You may choose the training codebase that best fits your requirements as long as you save the following files:
+1. The pretrained transformer layers (post-fine tuning) **separately** from the NER heads. file should be named `pytorch_model.bin`
+2. All the files required to be loaded by `transformers.AutoModel`, including the standard `config.json`, `special_tokens.json`, `tokenizer_config.json`, `vocab.txt`.
+3. Each of the NER head weights should be saved as separate `bin` file with a corresponding index to class mapping `json` file (see instructions below). Files should follow the `ner_<name>.bin` and `ner_<name>.json` name convention.
+
+**All weights files should be trained using Hugging Face and PyTorch libraries**
+
+For example, the following directory contains all the required files for a model that was jointly trained against the BMC and NEMO corpora:
+````
+resources/
+├── config.json
+├── pytorch_model.bin
+├── special_tokens.json
+├── tokenizer_config.json
+├── vocab.txt
+├── ner_bmc.bin
+├── ner_bmc.json
+├── ner_nemo.bin
+└── ner_nemo.json
+````
+
+### Index to class mapping file
+Each NER head should include a `json` file that maps between the model class index to the corresponding token class name.
+Note that indices `0` and `1` should always be associated with `[PAD]` and `O` classes, respectively. Also, the token annotation schema should be [**IBO2**](https://en.wikipedia.org/wiki/Inside%E2%80%93outside%E2%80%93beginning_(tagging)).
+
+
+Here is an example of the index to class mapping `json` file for `ner_bmc`:
+````json
+{
+	"0": "[PAD]",
+	"1": "O",
+	"2": "B-PERS",
+	"3": "I-PERS",
+	"4": "B-LOC",
+	"5": "I-LOC",
+	"6": "B-ORG",
+	"7": "I-ORG",
+	"8": "B-TIME",
+	"9": "I-TIME",
+	"10": "B-DATE",
+	"11": "I-DATE",
+	"12": "B-MONEY",
+	"13": "I-MONEY",
+	"14": "B-PERCENT",
+	"15": "I-PERCENT",
+	"16": "B-MISC__AFF",
+	"17": "I-MISC__AFF",
+	"18": "B-MISC__ENT",
+	"19": "I-MISC__ENT",
+	"20": "B-MISC_EVENT",
+	"21": "I-MISC_EVENT"
+}
+````
+
+## Model Packaging
+Once you have prepared all the directory with all required files, please follow these steps:
+1. Fork this repo (in case you haven't already)
+2. Make sure that `spacy` is installed in your running python environment (**make sure it is the same version as mentioned in requirements.txt**) 
+3. Navigate to the repo's root
+4. Run `python setup.py develop`, which should create a `hebspacy.egg-info` directory
+5. Navigate to `scripts\model`
+6. Update `meta.json` accordingly (make sure to follow the [spaCy package naming conventions](https://spacy.io/models#conventions))
+7. Run `python package.py <RESOURCES_DIR> <OUTPUT_DIR>`, where `<RESOURCES_DIR>` should point to the directory with all the files from the previous section.
+8. Run `python -m spacy package <OUTPUT_DIR> <WHEEL_DIR> --build wheel`
+9. Your `whl` and `tar.gz` files are ready under `<WHEEL_DIR>/<lang>_<name>-<version>/dist`
+10. Install your files by running `pip install XXXXX.whl`
\ No newline at end of file
diff --git a/contribute/model/meta.json b/contribute/model/meta.json
new file mode 100644
index 0000000..2cd1797
--- /dev/null
+++ b/contribute/model/meta.json
@@ -0,0 +1,10 @@
+{
+  "lang": "he",
+  "name": "ner_news_trf",
+  "version": "3.2.1",
+  "description": "Hebrew pipeline. Components: tok2vec, senter, multihead ner (bmc and nemo).",
+  "author": "hebpsacy",
+  "email": "hebspacy@gmail.com",
+  "url": "https://github.com/8400TheHealthNetwork/HebSpacy",
+  "license": "MIT License (MIT)"
+}
\ No newline at end of file
diff --git a/contribute/model/package.py b/contribute/model/package.py
new file mode 100644
index 0000000..1ad11ed
--- /dev/null
+++ b/contribute/model/package.py
@@ -0,0 +1,55 @@
+import json
+import os
+import shutil
+from pathlib import Path
+from spacy.lang.xx import MultiLanguage
+from thinc.config import Config
+import argparse
+
+
+def package(resource_dir: str, output_dir: str):
+    input_dir = Path(resource_dir)
+    target_dir = Path(output_dir)
+
+    if target_dir.exists():
+        shutil.rmtree(target_dir)
+    target_dir.mkdir()
+
+    nlp = MultiLanguage()
+    # update metadata
+    with open("meta.json", "r") as f:
+        meta = json.load(f)
+    nlp.meta.update(meta)
+    nlp.lang = "he"
+
+    # Add sentence segmentation
+    nlp.add_pipe("sentencizer")
+
+    # Add transformer
+    config = Config().from_disk("transformer.cfg")
+    transformer_config = config["transformer"]
+    transformer_config["model"]["name"] = str(input_dir)
+    transformer = nlp.add_pipe("transformer", config=config["transformer"])
+    transformer.model.initialize()
+
+    # Add ner heads
+    for ner_head in input_dir.glob("ner_*.bin"):
+        pipe = nlp.add_pipe("ner_head", name=ner_head.stem, last=True)
+        pipe.load(input_dir)
+        os.mkdir(target_dir / ner_head.stem)
+
+    # Add consolidator
+    nlp.add_pipe("consolidator", last=True)
+
+    nlp.to_disk(target_dir)
+
+
+if __name__ == '__main__':
+    # make sure to run `python setup.py develop` before running this script
+    parser = argparse.ArgumentParser()
+    parser.add_argument('resources_dir', help="directory containing the weights and configuration files")
+    parser.add_argument('output_dir', default=".", help="pipeline output directory")
+
+    args = parser.parse_args()
+
+    package(args.resources_dir, args.output_dir)
diff --git a/contribute/model/transformer.cfg b/contribute/model/transformer.cfg
new file mode 100644
index 0000000..e21b9a0
--- /dev/null
+++ b/contribute/model/transformer.cfg
@@ -0,0 +1,16 @@
+[transformer]
+max_batch_items = 4096
+
+[transformer.set_extra_annotations]
+@annotation_setters = "spacy-transformers.null_annotation_setter.v1"
+
+[transformer.model]
+@architectures = "spacy-transformers.TransformerModel.v3"
+# name = WILL BE AUTOMATICALLY ADDED BY SCRIPT
+tokenizer_config = {"use_fast": true}
+transformer_config = {}
+mixed_precision = false
+grad_scaler_config = {}
+
+[transformer.model.get_spans]
+@span_getters = "spacy-transformers.sent_spans.v1"
diff --git a/contribute/pipeline/README.md b/contribute/pipeline/README.md
new file mode 100644
index 0000000..e551998
--- /dev/null
+++ b/contribute/pipeline/README.md
@@ -0,0 +1,10 @@
+# Contribute to `hebspacy` pipeline
+In case you added new feature, fixed bugs or made any changes in the repo, please follow these instructions:
+
+1. Run existing and new tests on your code. make sure to place them under `/tests`
+2. Navigate to `hebspacy/version.py` and promote the package version according to [semantic versioning scheme](https://packaging.python.org/en/latest/guides/distributing-packages-using-setuptools/#semantic-versioning-preferred)
+3. Run `python setup.py bdist_wheel` to create a new `whl` file
+4. You can now install your version by running `pip install XXXXX.whl`
+5. Once you are ready to upload your version to [PyPi](https://pypi.org/project/hebspacy/), please contact the package maintainers personally or via email to [hebspacy@gmail.com](mailto:hebspacy@gmail.com)
+
+Note that previous model files may need to be repackaged to work with the new version. 
\ No newline at end of file
diff --git a/hebspacy/__init__.py b/hebspacy/__init__.py
new file mode 100644
index 0000000..99c4176
--- /dev/null
+++ b/hebspacy/__init__.py
@@ -0,0 +1,1 @@
+__version__ = '0.0.1'
\ No newline at end of file
diff --git a/hebspacy/ner_head.py b/hebspacy/ner_head.py
new file mode 100644
index 0000000..da68067
--- /dev/null
+++ b/hebspacy/ner_head.py
@@ -0,0 +1,142 @@
+import inspect
+import json
+from collections import defaultdict
+
+import numpy as np
+import torch
+from spacy import Language
+from spacy.tokens import Doc, Span
+from spacy.training import iob_to_biluo, biluo_tags_to_spans
+from torch.nn import functional as funct
+
+
+@Language.factory("ner_head")
+class NERHead:
+    def __init__(self, nlp: Language, name: str):
+        self.name = name
+        Doc.set_extension(f"{name}_ents", default=None, force=True)
+        Span.set_extension("confidence_score", default=1.0, force=True)
+
+
+    def load(self, dir_path: str):
+        self.model = torch.load(f"{dir_path}/{self.name}.bin")
+        with open(f"{dir_path}/{self.name}.json", "r") as f:
+            self.idx_to_tag = json.load(f)
+
+    def __call__(self, doc):
+        tokens_pred = []
+        probs_pred = []
+        # sequence length is defined by the longest sentence (shorter sentences are padded sentences)
+        seq_size = len(doc._.trf_data.wordpieces.strings[-1])
+
+        for index, token in enumerate(doc):
+            # check if a word piece exist for the token
+            if len(doc._.trf_data.align[index].dataXd) > 0:
+                # retrieve the index of the token head word piece
+                wp_index = doc._.trf_data.align[index].dataXd[0][0]
+                # calculate the location of the vector when considering the batches
+                batch = wp_index // seq_size
+                item = wp_index % seq_size
+
+                # retrieving the vectors of the token's head word piece from the last 4 layers
+                layer_9 = torch.from_numpy(doc._.trf_data.tensors[2][8][batch][item])
+                layer_10 = torch.from_numpy(doc._.trf_data.tensors[2][9][batch][item])
+                layer_11 = torch.from_numpy(doc._.trf_data.tensors[2][10][batch][item])
+                layer_12 = torch.from_numpy(doc._.trf_data.tensors[2][11][batch][item])
+
+                # concat the vectors into one
+                concatenated_vector = torch.cat([layer_9, layer_10, layer_11, layer_12])
+
+                # calculate the logits
+                logits = funct.linear(concatenated_vector, self.model["weight"], self.model["bias"])
+
+                pred = torch.argmax(logits).numpy().tolist()
+
+                # calculate probabilities and save the score for the top class
+                prob = torch.softmax(torch.FloatTensor(logits), dim=-1)
+                probs_pred.append(prob[pred].tolist())
+
+                if pred == 0:
+                    pred = 1
+            else:
+                pred = 1
+                probs_pred.append(1)
+            tokens_pred.append(self.idx_to_tag[str(pred)])
+        biluo_tags = iob_to_biluo(tokens_pred)
+        spans = biluo_tags_to_spans(doc, biluo_tags)
+
+        # entity confidence score is the average score across the corresponding tokens
+        for span in spans:
+            span._.confidence_score = np.average(probs_pred[span.start:span.end])
+        doc = self.set_entities(doc, spans)
+
+        return doc
+
+    def to_disk(self, path, exclude=tuple()):
+        # This will receive the directory path + /my_component
+        model_path = path / "model.bin"
+        torch.save(self.model, model_path)
+        idx_to_tag_mapping_path = path / "idx_to_tag.json"
+        with idx_to_tag_mapping_path.open("w", encoding="utf8") as f:
+            f.write(json.dumps(self.idx_to_tag))
+
+    def from_disk(self, path, exclude=tuple()):
+        # This will receive the directory path + /my_component
+        model_path = path / "model.bin"
+        self.model = torch.load(model_path)
+        idx_to_tag_mapping_path = path / "idx_to_tag.json"
+        with idx_to_tag_mapping_path.open("r", encoding="utf8") as f:
+            self.idx_to_tag = json.load(f)
+        return self
+
+    def set_entities(self, doc, spans):
+        setattr(doc._, f"{self.name}_ents", spans)
+        return doc
+
+
+@Language.component("consolidator")
+def consolidator(doc):
+    """
+    Weights the recognized entities from the different ner heads and selects a subset that do not non-overlap
+    The current consolidation policy prefers entities first by starting position and then by length
+    """
+
+    # extract all all ner_head extensions
+    extensions = filter(lambda ext: ext[0].endswith("_ents"), inspect.getmembers(doc._))
+    ents = []
+    start_to_entities = defaultdict(list)
+
+    # map start index to entities
+    for name, entities in extensions:
+        if entities:
+            for entity in entities:
+                start_to_entities[entity.start].append(entity)
+
+    # iterate over the starting position and select the longest entity
+    entities_start_positions = sorted(list(start_to_entities.keys()))
+    if len(entities_start_positions) > 0:
+        max_index = entities_start_positions[-1]
+        index = 0
+        while entities_start_positions[index] <= max_index:
+            start_index = entities_start_positions[index]
+            entities = start_to_entities[start_index]
+
+            # find the longest entity
+            ent_lens = [len(ent) for ent in entities]
+            max_len = max(ent_lens)
+            arg_max = ent_lens.index(max_len)
+
+            # select that entity
+            target_entity = entities[arg_max]
+            ents.append(target_entity)
+
+            # find the next index to investigate
+            min_target_index = start_index + len(target_entity)
+            if max_index < min_target_index:
+                break
+            while entities_start_positions[index] < min_target_index:
+                index += 1
+
+        doc.set_ents(ents)
+    return doc
+
diff --git a/hebspacy/version.py b/hebspacy/version.py
new file mode 100644
index 0000000..45d6d0d
--- /dev/null
+++ b/hebspacy/version.py
@@ -0,0 +1,5 @@
+_MAJOR = "0"
+_MINOR = "1"
+_REVISION = "7"
+
+VERSION = f"{_MAJOR}.{_MINOR}.{_REVISION}"
diff --git a/requirements-dev.txt b/requirements-dev.txt
new file mode 100644
index 0000000..55b033e
--- /dev/null
+++ b/requirements-dev.txt
@@ -0,0 +1,1 @@
+pytest
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..f135a7a
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,6 @@
+sentencepiece==0.1.96
+spacy==3.2.2
+spacy-transformers
+tokenizers
+transformers
+torch
\ No newline at end of file
diff --git a/setup.cfg b/setup.cfg
new file mode 100644
index 0000000..3287428
--- /dev/null
+++ b/setup.cfg
@@ -0,0 +1,3 @@
+# Inside of setup.cfg
+[metadata]
+description-file = README.md
diff --git a/setup.py b/setup.py
new file mode 100644
index 0000000..6192c27
--- /dev/null
+++ b/setup.py
@@ -0,0 +1,64 @@
+import os
+from os import path
+
+from setuptools import setup, find_packages
+
+this_directory = path.abspath(path.dirname(__file__))
+parent_directory = os.path.abspath(os.path.join(this_directory, os.pardir))
+
+VERSION = {}
+with open(f"{this_directory}/hebspacy/version.py", "r") as version_file:
+    exec(version_file.read(), VERSION)
+
+def get_requirements_from_files(*files):
+    install_requires = {}
+    for filename in files:
+        if os.path.isfile(filename):
+            with open(filename) as f:
+                for line in f.readlines():
+                    line = line.strip()
+                    if line and not line.startswith("#"):
+                        if "#" in line:
+                            req, version = line.split("#")[0].strip().split("==")
+                        else:
+                            req, version = line, None
+
+                        # requirement versions must be specified with "==", otherwise the line above will fail
+                        if req not in install_requires:
+                            install_requires[req] = version
+    reqs = []
+    for k, v in sorted(install_requires.items()):
+        if v:
+            reqs.append(f"{k}=={v}")
+        else:
+            reqs.append(k)
+    return reqs
+
+
+install_requires_ = get_requirements_from_files('requirements.txt')
+
+for requirement in install_requires_:
+    print("adding requirement: " + requirement)
+
+with open(path.join(this_directory, "README.MD"), encoding="utf-8") as f:
+    long_description = f.read()
+
+setup(
+    name="hebspacy",
+    author="hebpsacy",
+    version=VERSION["VERSION"],
+    author_email="hebspacy@gmail.com",
+    url="https://github.com/8400TheHealthNetwork/HebSpacy",
+    license="MIT License (MIT)",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    description="SpaCy pipeline and models for Hebrew text",
+    keywords=["hebrew nlp spacy SpaCy phi pii"],
+    entry_points={
+        "spacy_factories": ["ner_head = hebspacy.ner_head:NERHead",
+                            "consolidator = hebspacy.ner_head:consolidator"]
+    },
+    packages=find_packages(),
+    install_requires=install_requires_,
+    python_requires=">=3.8.0",
+)


Test patch:
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 0000000..30bc8cb
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,7 @@
+import pytest
+import spacy
+
+
+@pytest.fixture(scope="session")
+def he_vocab():
+    return spacy.util.get_lang_class("he")().vocab
diff --git a/tests/test_consolidator.py b/tests/test_consolidator.py
new file mode 100644
index 0000000..9fa05be
--- /dev/null
+++ b/tests/test_consolidator.py
@@ -0,0 +1,100 @@
+import pytest
+
+from hebspacy.ner_head import NERHead, consolidator
+from spacy.tokens import Doc, Span
+
+
+def test_consolidator_with_single_head(he_vocab):
+    """
+    Checks that ner head entities are populated to ents property post consolidation
+    """
+    mock_doc = Doc(he_vocab, words=["שלום", "כיתה", "אלף"])
+    new_head = NERHead(nlp=None, name="test")
+    entities = [Span(mock_doc, 0, 1, label="PERSON"), Span(mock_doc, 1, 2, label="LOC")]
+    new_head.set_entities(mock_doc, entities)
+    doc = consolidator(mock_doc)
+
+    assert doc == mock_doc
+    assert len(doc.ents) == 2
+    assert doc.ents[0] == entities[0]
+    assert doc.ents[1] == entities[1]
+
+
+def test_consolidator_with_multiple_heads_ent_contained_in_ent(he_vocab):
+    """
+    Checks that in case of entity conflict, the consolidator prefers the longer term (assuming they share the same index)
+    """
+    mock_doc = Doc(he_vocab, words=["שלום", "כיתה", "אלף"])
+
+    new_head = NERHead(nlp=None, name="test")
+    entities_test = [Span(mock_doc, 0, 1, label="PERSON"), Span(mock_doc, 1, 2, label="LOC")]
+    new_head.set_entities(mock_doc, entities_test)
+
+    new_head = NERHead(nlp=None, name="test_2")
+    entities_test_2 = [Span(mock_doc, 0, 1, label="PERSON"), Span(mock_doc, 1, 3, label="LOC")]
+    new_head.set_entities(mock_doc, entities_test_2)
+
+    doc = consolidator(mock_doc)
+
+    assert doc == mock_doc
+    assert len(doc.ents) == 2
+    assert doc.ents[0].start == entities_test[0].start and doc.ents[0].end == entities_test[0].end
+    assert doc.ents[1].start == entities_test_2[1].start and doc.ents[1].end == entities_test_2[1].end
+
+
+def test_consolidator_with_multiple_heads_ents_conflicts(he_vocab):
+    """
+    Checks that in case of entity conflict/paritial overlap, the consolidator prefers position
+    """
+    mock_doc = Doc(he_vocab, words=["שלום", "כיתה", "אלף", "8"])
+
+    new_head = NERHead(nlp=None, name="test")
+    entities_test = [Span(mock_doc, 0, 1, label="PERSON"), Span(mock_doc, 1, 3, label="LOC")]
+    new_head.set_entities(mock_doc, entities_test)
+
+    new_head = NERHead(nlp=None, name="test_2")
+    entities_test_2 = [Span(mock_doc, 0, 1, label="PERSON"), Span(mock_doc, 2, 4, label="LOC")]
+    new_head.set_entities(mock_doc, entities_test_2)
+
+    doc = consolidator(mock_doc)
+
+    assert doc == mock_doc
+    assert len(doc.ents) == 2
+    assert doc.ents[0].start == entities_test[0].start and doc.ents[0].end == entities_test[0].end
+    assert doc.ents[1].start == entities_test[1].start and doc.ents[1].end == entities_test[1].end
+
+
+def test_consolidator_with_no_heads(he_vocab):
+    """
+    make sure consolidator doesn't crash in case there are no ner heads
+    """
+    mock_doc = Doc(he_vocab, words=["שלום", "כיתה", "אלף"])
+    try:
+        doc = consolidator(mock_doc)
+    except Exception as e:
+        pytest.fail("Unexpected error", str(e))
+
+
+def test_consolidator_with_head_not_initialized(he_vocab):
+    """
+    make sure consolidator doesn't crash in case ner heads weren't initialized
+    """
+    mock_doc = Doc(he_vocab, words=["שלום", "כיתה", "אלף"])
+    new_head = NERHead(nlp=None, name="test")
+    try:
+        doc = consolidator(mock_doc)
+    except IndexError as e:
+        pytest.fail("Unexpected Index error", str(e))
+
+
+def test_consolidator_with_head_no_entities(he_vocab):
+    """
+    make sure consolidator doesn't crash in case no entities were detected
+    """
+    mock_doc = Doc(he_vocab, words=["שלום", "כיתה", "אלף"])
+    new_head = NERHead(nlp=None, name="test")
+    new_head.set_entities(mock_doc, [])
+    try:
+        doc = consolidator(mock_doc)
+    except IndexError as e:
+        pytest.fail("Unexpected Index error", str(e))
diff --git a/tests/test_ner_head.py b/tests/test_ner_head.py
new file mode 100644
index 0000000..20ae2a7
--- /dev/null
+++ b/tests/test_ner_head.py
@@ -0,0 +1,57 @@
+import pytest
+from hebspacy.ner_head import NERHead
+from spacy.tokens import Doc, Span
+from spacy.tokens.doc import Underscore
+
+
+@pytest.fixture(autouse=True)
+def remove_extensions():
+    """
+    Remove extensions as it is global
+    """
+    if Span.has_extension("confidence_score"):
+        Span.remove_extension("confidence_score")
+
+    # remove ner head entities extensions ({name}_ents)
+    doc_extension_names = list(filter(lambda name: name.endswith("_ents"), Underscore.doc_extensions.keys()))
+    for name in doc_extension_names:
+        Underscore.doc_extensions.pop(name)
+
+
+def test_span_confidence_score_extension_added(he_vocab):
+    """
+    Check that confidence_score extension is available
+    """
+
+    new_head = NERHead(nlp=None, name="test")
+    mock_doc = Doc(he_vocab, words=["שלום", "כיתה", "אלף"])
+
+    assert Span.has_extension("confidence_score")
+
+
+def test_confidence_score(he_vocab):
+    """
+    Check that confidence_score works is properly set
+    """
+    new_head = NERHead(nlp=None, name="test")
+    mock_doc = Doc(he_vocab, words=["שלום", "כיתה", "אלף"])
+    span = Span(mock_doc, 0, 1, label="PERSON")
+    span._.confidence_score = 0.92
+
+    new_head.set_entities(mock_doc, [span])
+
+    assert mock_doc._.test_ents[0] == span
+    assert mock_doc._.test_ents[0]._.confidence_score == span._.confidence_score
+
+
+def test_create_new_head_and_set_entities(he_vocab):
+    """
+    Check that creating a new head generates a corresponding extension with the entities
+    """
+    new_head = NERHead(nlp=None, name="test")
+    mock_doc = Doc(he_vocab, words=["שלום", "כיתה", "אלף"])
+    entities = [Span(mock_doc, 0, 1, label="PERSON"), Span(mock_doc, 1, 2, label="LOC")]
+    new_head.set_entities(mock_doc, entities)
+
+    assert Doc.has_extension("test_ents")
+    assert mock_doc._.test_ents == entities