Repository: transformers
Issue: Feature Request: To add nested hierarchy retrieval from Donut response
Labels: Feature request

### Feature request

### Donut for hierarchy extraction (Document Parsing)
While preprocessing the ground truth json to the tokens for Donut the processor function (json2token) handles nested hierarchy but the same doesn't hold true for token2json.

Below is an example json:
`
{
    "header": "This is 1st header",
    "elements": [
      {
        "text_block": "This is a textblock"
      },
      {
        "header": "1st nested header",
          "elements": [
            {
              "text_block": "This is a sentence"
            },
            {
              "text_block": "Another sentence...."
            },
            {
              "itallic_header": "This is an itallic header",
              "elements": [
                {
                  "text_block": "Text 1 inside itallic header.."
                },
                {
                  "text_block": "Text 2 inside itallic header.."
                }
              ]
            }
        ]
      }
    ]
  }
`

Consider the above json. Applying the json2token function gives the following token sequence.

Function Call:
`output = json2token(temp_test)`

> <s_header>This is 1st header</s_header><s_elements><s_text_block>This is a textblock</s_text_block><sep/><s_header>1st nested header</s_header><s_elements><s_text_block>This is a sentence</s_text_block><sep/><s_text_block>Another sentence....</s_text_block><sep/><s_itallic_header>This is an itallic header</s_itallic_header><s_elements><s_text_block>Text 1 inside itallic header..</s_text_block><sep/><s_text_block>Text 2 inside itallic header..</s_text_block></s_elements></s_elements></s_elements>

This maintains the hierarchy (like parenthesis matching). So, if donut is trained on such data it will give response which parses the information  & also retains the hierarchy but the token2json function doesn't handle the conversion properly. Below is the output of the function id passed the token sequence present above.

Function Call:
`processor.token2json(output)`

Output

`
[
  {
    'header': 'This is 1st header',
    'elements': [
      {
        'text_block': 'This is a textblock'
      },
      {
        'header': '1st nested header',
        'text_block': 'This is a sentence'
      },
      {
        'text_block': 'Another sentence....'
      },
      {
        'itallic_header': 'This is an itallic header',
        'text_block': 'Text 1 inside itallic header..'
      },
      {
        'text_block': 'Text 2 inside itallic header..'
      }
    ]
  }
]
`

Updated Function Results (Preserving the hierarchy):
`
[
  {
    'header': 'This is 1st header',
    'elements': [
      {
        'text_block': 'This is a textblock'
      },
      {
        'header': '1st nested header',
        'elements': [
          {
            'text_block': 'This is a sentence'
          },
          {
            'text_block': 'Another sentence....'
          },
          {
            'itallic_header': 'This is an itallic header',
            'elements': [
              {
                'text_block': 'Text 1 inside itallic header..'
              },
              {
                'text_block': 'Text 2 inside itallic header..'
              }
            ]
          }
        ]
      }
    ]
  }
]
`

Example from CORD:

> temp_test = {
  "company": "ADVANCO COMPANY",
  "date": "17/01/2018",
  "address": "NO 1&3, JALAN WANGSA DELIMA 12, WANGSA LINK, WANGSA MAJU, 53300 KUALA LUMPUR",
  "total": "7.00"
}

Updated Function Output:
`
[
  {
    'company': 'ADVANCO COMPANY',
    'date': '17/01/2018',
    'address': 'NO 1&3, JALAN WANGSA DELIMA 12, WANGSA LINK, WANGSA MAJU, 53300 KUALA LUMPUR',
    'total': '7.00'
  }
]
`

### Motivation

Found out about this while working on a project to extract information from images also maintaining the hierarchy/structure of it.

Going through the CORD dataset made me realize that the data itself is not nested in nature. So, thought of testing on a sample the postprocessing logics json -> token & token -> json conversion.

Updated the token2json to get the hierarchy as it is from the token but wasn't sure about the model performance on nested jsons but long story short Donut predicts the hierarchy pretty good.

### Your contribution

`
def token2json(tokens, is_inner_value=False, nested_key = 'elements'):
    """
    Convert a (generated) token seuqnce into an ordered JSON format
    """
    output = dict()

    while tokens:
        start_token = re.search(r"<s_(.*?)>", tokens, re.IGNORECASE)
        if start_token is None:
            break
        key = start_token.group(1)
        start_matches = re.finditer(fr"<s_{key}>", tokens)
        end_matches = re.finditer(fr"</s_{key}>", tokens)
        start_tups = [(match.group(), match.start(), match.end()) for match in start_matches]
        end_tups = [(match.group(), match.start(), match.end()) for match in end_matches]
        mergeTups = start_tups + end_tups
        sortedMergeTups = sorted(mergeTups, key=lambda x: x[1])
        # remove any unattended close tag for the key present before the current focus start key
        updatedIdx = -1
        for idx in range(len(sortedMergeTups)):
            if start_token.span()[0] == sortedMergeTups[idx][1]:
                updatedIdx = idx
                break
        sortedMergeTups = sortedMergeTups[updatedIdx:]
        start_main = sortedMergeTups[0]
        match_tracker = 0
        end_token = None
        if key == nested_key :
            if start_main[0] == f'<s_{key}>':
                for tup in sortedMergeTups[1:]:
                    if tup[0] == f'</s_{key}>':
                        if match_tracker == 0:
                            end_token = tup
                            break
                        else:
                            match_tracker -= 1
                    elif tup[0] == f'<s_{key}>':
                        match_tracker += 1
        elif len(sortedMergeTups) > 1:
            nextTup = sortedMergeTups[1]
            if nextTup[0] == f'</s_{key}>':
                end_token = nextTup

        if end_token is None:
            tokens = tokens.replace(start_token[0], "", 1)
        else:
            start_token_word = start_main[0]
            start_token_id = start_main[2]
            end_token_word = end_token[0]
            end_token_id = end_token[1]
            content = tokens[start_token_id: end_token_id]
            if content is not None:
                if r"<s_" in content and r"</s_" in content:  # non-leaf node
                    value = token2json(content, is_inner_value=True)
                    if value:
                        if len(value) == 1:
                            value = value[0]
                        output[key] = value
                else:  # leaf nodes
                    if key in output.keys():
                        if isinstance(output[key], str):
                            tempVal = output[key]
                            output[key] = [tempVal]
                    else:    
                        output[key] = []
                    for leaf in content.split(r"<sep/>"):
                        leaf = leaf.strip()
                        if (
                            leaf in processor.tokenizer.get_added_vocab()
                            and leaf[0] == "<"
                            and leaf[-2:] == "/>"
                        ):
                            leaf = leaf[1:-2]  # for categorical special tokens
                        output[key].append(leaf)
                    if len(output[key]) == 1:
                        output[key] = output[key][0]

            tokens = tokens[end_token[2]:]
            if tokens[:6] == r"<sep/>":  # non-leaf nodes
                return [output] + token2json(tokens[6:], is_inner_value=True)

    if len(output):
        return [output] if is_inner_value else output
    else:
        return [] if is_inner_value else {"text_sequence": tokens}
`