Repository: flairNLP/flair

Issue:
Enhancement: NER Output not as string but nested dict
It would be great if there was a way to obtain a nested dictionary as output of the NER instead of a string with <...> tags. The string is quite tedious to work with.

I imagine an output like:

```
"sentence": {
    "text": "Facebook, Inc. is a company, and Google is one as well.",
    "named_entities": [
        {
            "mention_text": "Facebook, Inc.",
            "start_pos": 0
            "end_pos": ...
            "type": "ORG",
            "confidence": 0.9
        },
        {
        ...
        }
    ]
} 
```


Discussion:
Hi, that's a very good point! 

We are planning to add the confidence of each prediction to the output (#38), and also the token offsets, for which we want to add the 'space_after' field like in the CoNLL-U format to remember after tokenization the original whitespacing of the sentence. This is very likely to be in the next release! 

In the meantime, you can also work with the Token objects instead of the string with the <..> tags. You can iterate through the tokens in a sentence and directly access their tag fields: 

```python
for token in sentence:
    print(token)
    print(token.get_tag('ner'))
```

Hope this helps!



Resolving patch:
diff --git a/flair/data.py b/flair/data.py
--- a/flair/data.py
+++ b/flair/data.py
@@ -129,6 +129,12 @@ def score(self, score):
         else:
             self._score = 1.0
 
+    def to_dict(self):
+        return {
+            'value': self.value,
+            'confidence': self.score
+        }
+
     def __str__(self):
         return "{} ({})".format(self._value, self._score)
 
@@ -147,12 +153,16 @@ def __init__(self,
                  idx: int = None,
                  head_id: int = None,
                  whitespace_after: bool = True,
+                 start_position: int = None
                  ):
         self.text: str = text
         self.idx: int = idx
         self.head_id: int = head_id
         self.whitespace_after: bool = whitespace_after
 
+        self.start_pos = start_position
+        self.end_pos = start_position + len(text) if start_position is not None else None
+
         self.sentence: Sentence = None
         self._embeddings: Dict = {}
         self.tags: Dict[str, Label] = {}
@@ -168,12 +178,6 @@ def get_tag(self, tag_type: str) -> Label:
     def get_head(self):
         return self.sentence.get_token(self.head_id)
 
-    def __str__(self) -> str:
-        return 'Token: %d %s' % (self.idx, self.text) if self.idx is not None else 'Token: %s' % (self.text)
-
-    def __repr__(self) -> str:
-        return 'Token: %d %s' % (self.idx, self.text) if self.idx is not None else 'Token: %s' % (self.text)
-
     def set_embedding(self, name: str, vector: torch.autograd.Variable):
         self._embeddings[name] = vector.cpu()
 
@@ -189,10 +193,24 @@ def get_embedding(self) -> torch.FloatTensor:
 
         return torch.FloatTensor()
 
+    @property
+    def start_position(self) -> int:
+        return self.start_pos
+
+    @property
+    def end_position(self) -> int:
+        return self.end_pos
+
     @property
     def embedding(self):
         return self.get_embedding()
 
+    def __str__(self) -> str:
+        return 'Token: %d %s' % (self.idx, self.text) if self.idx is not None else 'Token: %s' % (self.text)
+
+    def __repr__(self) -> str:
+        return 'Token: %d %s' % (self.idx, self.text) if self.idx is not None else 'Token: %s' % (self.text)
+
 
 class Span:
     """
@@ -203,11 +221,39 @@ def __init__(self, tokens: List[Token], tag: str = None, score=1.):
         self.tokens = tokens
         self.tag = tag
         self.score = score
+        self.start_pos = None
+        self.end_pos = None
+
+        if tokens:
+            self.start_pos = tokens[0].start_position
+            self.end_pos = tokens[len(tokens) - 1].end_position
 
     @property
     def text(self) -> str:
         return ' '.join([t.text for t in self.tokens])
 
+    def to_original_text(self) -> str:
+        str = ''
+        pos = self.tokens[0].start_pos
+        for t in self.tokens:
+            while t.start_pos != pos:
+                str += ' '
+                pos += 1
+
+            str += t.text
+            pos += len(t.text)
+
+        return str
+
+    def to_dict(self):
+        return {
+            'text': self.to_original_text(),
+            'start_pos': self.start_pos,
+            'end_pos': self.end_pos,
+            'type': self.tag,
+            'confidence': self.score
+        }
+
     def __str__(self) -> str:
         ids = ','.join([str(t.idx) for t in self.tokens])
         return '{}-span [{}]: "{}"'.format(self.tag, ids, self.text) \
@@ -254,7 +300,7 @@ def __init__(self, text: str = None, use_tokenizer: bool = False, labels: Union[
                 last_word_offset = -1
                 last_token = None
                 for word in tokens:
-                    token = Token(word)
+                    token = Token(word, start_position=index(word, running_offset))
                     self.add_token(token)
                     try:
                         word_offset = index(word, running_offset)
@@ -270,10 +316,12 @@ def __init__(self, text: str = None, use_tokenizer: bool = False, labels: Union[
             # otherwise assumes whitespace tokenized text
             else:
                 # add each word in tokenized string as Token object to Sentence
+                offset = 0
                 for word in text.split(' '):
                     if word:
-                        token = Token(word)
+                        token = Token(word, start_position=text.index(word, offset))
                         self.add_token(token)
+                        offset += len(word) + 1
 
     def get_token(self, token_id: int) -> Token:
         for token in self.tokens:
@@ -473,6 +521,34 @@ def infer_space_after(self):
             last_token = token
         return self
 
+    def to_original_text(self) -> str:
+        str = ''
+        pos = 0
+        for t in self.tokens:
+            while t.start_pos != pos:
+                str += ' '
+                pos += 1
+
+            str += t.text
+            pos += len(t.text)
+
+        return str
+
+    def to_dict(self, tag_type: str = None):
+        labels = []
+        entities = []
+
+        if tag_type:
+            entities = [span.to_dict() for span in self.get_spans(tag_type)]
+        if self.labels:
+            labels = [l.to_dict() for l in self.labels]
+
+        return {
+            'text': self.to_original_text(),
+            'labels': labels,
+            'entities': entities
+        }
+
     def __getitem__(self, idx: int) -> Token:
         return self.tokens[idx]
 


Test patch:
diff --git a/tests/test_data.py b/tests/test_data.py
--- a/tests/test_data.py
+++ b/tests/test_data.py
@@ -448,4 +448,57 @@ def test_spans():
     assert (1 == len(spans))
 
     spans: List[Span] = sentence.get_spans('ner', min_score=0.99)
-    assert (0 == len(spans))
\ No newline at end of file
+    assert (0 == len(spans))
+
+
+def test_token_position_in_sentence():
+    sentence = Sentence("I love Berlin .")
+
+    assert(0 == sentence.tokens[0].start_position)
+    assert(1 == sentence.tokens[0].end_position)
+    assert(2 == sentence.tokens[1].start_position)
+    assert(6 == sentence.tokens[1].end_position)
+    assert(7 == sentence.tokens[2].start_position)
+    assert(13 == sentence.tokens[2].end_position)
+
+    sentence = Sentence(" I love  Berlin.", use_tokenizer=True)
+
+    assert(1 == sentence.tokens[0].start_position)
+    assert(2 == sentence.tokens[0].end_position)
+    assert(3 == sentence.tokens[1].start_position)
+    assert(7 == sentence.tokens[1].end_position)
+    assert(9 == sentence.tokens[2].start_position)
+    assert(15 == sentence.tokens[2].end_position)
+
+
+def test_sentence_to_dict():
+    sentence = Sentence('Zalando Research is   located in Berlin, the capital of Germany.', labels=['business'], use_tokenizer=True)
+
+    # bioes tags
+    sentence[0].add_tag('ner', 'B-ORG')
+    sentence[1].add_tag('ner', 'E-ORG')
+    sentence[5].add_tag('ner', 'S-LOC')
+    sentence[10].add_tag('ner', 'S-LOC')
+
+    dict = sentence.to_dict('ner')
+
+    assert ('Zalando Research is   located in Berlin, the capital of Germany.' == dict['text'])
+    assert ('Zalando Research' == dict['entities'][0]['text'])
+    assert ('Berlin' == dict['entities'][1]['text'])
+    assert ('Germany' == dict['entities'][2]['text'])
+    assert (1 == len(dict['labels']))
+
+    sentence = Sentence('Facebook, Inc. is a company, and Google is one as well.', use_tokenizer=True)
+
+    # bioes tags
+    sentence[0].add_tag('ner', 'B-ORG')
+    sentence[1].add_tag('ner', 'I-ORG')
+    sentence[2].add_tag('ner', 'E-ORG')
+    sentence[8].add_tag('ner', 'S-ORG')
+
+    dict = sentence.to_dict('ner')
+
+    assert ('Facebook, Inc. is a company, and Google is one as well.' == dict['text'])
+    assert ('Facebook, Inc.' == dict['entities'][0]['text'])
+    assert ('Google' == dict['entities'][1]['text'])
+    assert (0 == len(dict['labels']))
\ No newline at end of file