Convert NER prediction to BIO format

Closes #229 (closed)

Tested with the provided Hennessy archive and the following script:

import logging
from pathlib import Path

from dan.bio import convert
from dan.utils import parse_tokens

logger = logging.getLogger(__name__)

ner_tokens = parse_tokens("data_hennessy/tokens.yml")

for directory in ["GT", "preds"]:
    base_output = Path(f"bio/{directory}")
    base_output.mkdir(parents=True, exist_ok=True)

    for path in Path(f"data_hennessy/{directory}").iterdir():
        logger.info(f"Converting file {path}")

        try:
            conversion = convert(path.read_text(), ner_tokens)
        except Exception as e:
            logger.error(e)
            continue

        output = base_output / path.name
        output.write_text(conversion)

Conversion on 2b325b53-0a27-4d6c-8325-0190e1c5c51b.txt

Text input

London Ⓐ20th January 1792Ⓑ
I wrote My Dear ⒸJimmyⒹ the 17th Instant and yesterday received yours of the 9th Ditto and Note Its contents am surprized my missing letters were not Come to hand, tI do not believe that ⒸBennet and HinshowⒹ had wrote you, I mentioned to you in my last that tho there was a Mistake in their order to me for the 20 Puncheons they would accept of them at the price you would have Purchased on the 7th Instant I mentioned some posts ago that Messieurs ⒸKneyon and PeileⒹ also accepted of the Purchase you in suspence you are well covered should they leave them for y a which I do not Immagin will be the case, I told you that when Master ⒸScholeyⒹ spoak to me of your Proposal he did not tell me if he would accept of it or not, but said he would write, I have heared nothing from him since on the subject, I call on him to today he was out and not Expected untill 4 ⁇Clock, I was 3 times today at ⒸT. BrownsⒹ 3 times at ⒸBevanⒹ’s saw neither of them I saw all the others in this quarter Nothing new ⒸTuffinⒹ only talkes to me of the affairs of ⒸTranceⒹ could wish himself a french man to take an active part in support of Liberty which he has now no doubt but they will Enjoy how is it that you Conceive so Bad an oppinion of this affair, I see it in a Very different light I have scare a doubt of its success, sir ab… and ⒸWilsonⒹ in the Country the samples per ⒸWilsonⒹ please them, my last Visit this afternoon was at ⒸHarrison W. and JohnsonsⒹ

BIO output

London O
20th B-date de la lettre
January I-date de la lettre
1792 I-date de la lettre
I O
wrote O
My O
Dear O
Jimmy B-nom propre
the O
17th O
Instant O
and O
yesterday O
received O
yours O
of O
the O
9th O
Ditto O
and O
Note O
Its O
contents O
am O
surprized O
my O
missing O
letters O
were O
not O
Come O
to O
hand, O
tI O
do O
not O
believe O
that O
Bennet B-nom propre
and I-nom propre
Hinshow I-nom propre
had O
wrote O
you, O
I O
mentioned O
to O
you O
in O
my O
last O
that O
tho O
there O
was O
a O
Mistake O
in O
their O
order O
to O
me O
for O
the O
20 O
Puncheons O
they O
would O
accept O
of O
them O
at O
the O
price O
you O
would O
have O
Purchased O
on O
the O
7th O
Instant O
I O
mentioned O
some O
posts O
ago O
that O
Messieurs O
Kneyon B-nom propre
and I-nom propre
Peile I-nom propre
also O
accepted O
of O
the O
Purchase O
you O
in O
suspence O
you O
are O
well O
covered O
should O
they O
leave O
them O
for O
y O
a O
which O
I O
do O
not O
Immagin O
will O
be O
the O
case, O
I O
told O
you O
that O
when O
Master O
Scholey B-nom propre
spoak O
to O
me O
of O
your O
Proposal O
he O
did O
not O
tell O
me O
if O
he O
would O
accept O
of O
it O
or O
not, O
but O
said O
he O
would O
write, O
I O
have O
heared O
nothing O
from O
him O
since O
on O
the O
subject, O
I O
call O
on O
him O
to O
today O
he O
was O
out O
and O
not O
Expected O
untill O
4 O
⁇Clock, O
I O
was O
3 O
times O
today O
at O
T. B-nom propre
Browns I-nom propre
3 O
times O
at O
Bevan B-nom propre
’s O
saw O
neither O
of O
them O
I O
saw O
all O
the O
others O
in O
this O
quarter O
Nothing O
new O
Tuffin B-nom propre
only O
talkes O
to O
me O
of O
the O
affairs O
of O
Trance B-nom propre
could O
wish O
himself O
a O
french O
man O
to O
take O
an O
active O
part O
in O
support O
of O
Liberty O
which O
he O
has O
now O
no O
doubt O
but O
they O
will O
Enjoy O
how O
is O
it O
that O
you O
Conceive O
so O
Bad O
an O
oppinion O
of O
this O
affair, O
I O
see O
it O
in O
a O
Very O
different O
light O
I O
have O
scare O
a O
doubt O
of O
its O
success, O
sir O
ab… O
and O
Wilson B-nom propre
in O
the O
Country O
the O
samples O
per O
Wilson B-nom propre
please O
them, O
my O
last O
Visit O
this O
afternoon O
was O
at O
Harrison B-nom propre
W. I-nom propre
and I-nom propre
Johnsons I-nom propre

dan/bio.py

-from typing import Dict
+from operator import attrgetter
+from typing import Dict, List
 
 from dan.utils import EntityType
 
 
+def find_starting_token(ner_tokens, ending_token):
+    for entity_type in ner_tokens.values():
+        if entity_type.end == ending_token:
+            return entity_type.start
+
+
+def find_entity_name(ner_tokens, starting_token):
+    for name, entity_type in ner_tokens.items():
+        if entity_type.start == starting_token:
+            return name
+
+
 def convert(text: str, ner_tokens: Dict[str, EntityType]) -> str:
-    iob_string: str = ""  # Full IOB formatted string
-    entity_types: list[str] = []  # Encountered entity types
-    inside: bool = False  # Whether we are inside an entity
+    # Spacing starting tokens and ending tokens (if necessary)
+    starting_tokens: List[str] = list(map(attrgetter("start"), ner_tokens.values()))
+    for starting_token in starting_tokens:
+        text = text.replace(starting_token, f" {starting_token} ")
 
-    for word in text.split():
-        if not word:
-            continue
+    ending_tokens: List[str] = list(map(attrgetter("end"), ner_tokens.values()))
+    for ending_token in ending_tokens:
+        text = text.replace(ending_token, f" {ending_token} ")
 
-        token_iterator = iter(ner_tokens.items())
-        while (ner_token := next(token_iterator, None)) is not None:
-            name, entity_type = ner_token
+    has_ending_tokens: bool = bool(len(ending_tokens))
 
-            if word[0] == entity_type.start:
-                word = word[1:]
+    iob: List[str] = []  # List of IOB formatted strings
+    entity_types: List[str] = []  # Encountered entity types
+    inside: bool = False  # Whether we are inside an entity
+    for token in text.split():
+        # Encountering a starting token
+        if token in starting_tokens:
+            entity_types.append(token)
+
+            # Stopping any current entity type
+            if has_ending_tokens:
+                inside = False
 
-                entity_types.append(entity_type.start)
+            continue
 
-                if entity_type.end:
-                    inside = False
-                    continue
+        # Encountering an ending token
+        elif has_ending_tokens and token in ending_tokens:
+            # Making sure this ending token closes the current entity
+            assert entity_types[-1] == find_starting_token(ner_tokens, token)
 
-            elif entity_type.end and word[-1] == entity_type.end:
-                # Make sure the token is the closing of the current entity
-                assert entity_types[-1] == entity_type.start
+            # Removing the current entity from the queue as it is its end
+            entity_types.pop()
 
-                word = word[:-1]
+            # If there is still entities in the queue, we continue in the parent one
+            # Else, we are not in any entity anymore
+            inside = bool(entity_types)
 
-                # Remove from queue
-                entity_types.pop()
+            continue
 
-                # if there is no more entity, you remove inside
-                # else we continue parent entity
-                inside = bool(entity_types)
+        # The token is not part of an entity
+        if not entity_types:
+            iob.append(f"{token} O")
+            continue
 
-            # Not a NER token
+        # The token is part of at least one entity
+        entity_name = find_entity_name(ner_tokens, entity_types[-1])
 
-            # If there is no current entity type
-            if not entity_types:
-                iob_string += f"\n{word} O"
-                continue
+        if inside:
+            # Inside the same entity
+            iob.append(f"{token} I-{entity_name}")
+            continue
 
-            # There is at least one entity type
-            if inside:
-                # We are inside an entity
-                iob_string += f"\n{word} I-{name}"
-            else:
-                # We are starting an entity
-                inside = True
-                iob_string += f"\n{word} B-{name}"
+        # Starting a new entity
+        iob.append(f"{token} B-{entity_name}")
+        inside = True
 
-    print(iob_string)
+    # Concatenating all formatted iob strings
+    return "\n".join(iob)