pipeline/nltk/corpus/reader/conll.py

# Natural Language Toolkit: CONLL Corpus Reader
#
# Copyright (C) 2001-2023 NLTK Project
# Author: Steven Bird <[email protected]>
#         Edward Loper <[email protected]>
# URL: <https://www.nltk.org/>
# For license information, see LICENSE.TXT

"""
Read CoNLL-style chunk fileids.
"""

import textwrap

from nltk.corpus.reader.api import *
from nltk.corpus.reader.util import *
from nltk.tag import map_tag
from nltk.tree import Tree
from nltk.util import LazyConcatenation, LazyMap


class ConllCorpusReader(CorpusReader):
    """
    A corpus reader for CoNLL-style files.  These files consist of a
    series of sentences, separated by blank lines.  Each sentence is
    encoded using a table (or "grid") of values, where each line
    corresponds to a single word, and each column corresponds to an
    annotation type.  The set of columns used by CoNLL-style files can
    vary from corpus to corpus; the ``ConllCorpusReader`` constructor
    therefore takes an argument, ``columntypes``, which is used to
    specify the columns that are used by a given corpus. By default
    columns are split by consecutive whitespaces, with the
    ``separator`` argument you can set a string to split by (e.g.
    ``\'\t\'``).


    @todo: Add support for reading from corpora where different
        parallel files contain different columns.
    @todo: Possibly add caching of the grid corpus view?  This would
        allow the same grid view to be used by different data access
        methods (eg words() and parsed_sents() could both share the
        same grid corpus view object).
    @todo: Better support for -DOCSTART-.  Currently, we just ignore
        it, but it could be used to define methods that retrieve a
        document at a time (eg parsed_documents()).
    """

    # /////////////////////////////////////////////////////////////////
    # Column Types
    # /////////////////////////////////////////////////////////////////

    WORDS = "words"  #: column type for words
    POS = "pos"  #: column type for part-of-speech tags
    TREE = "tree"  #: column type for parse trees
    CHUNK = "chunk"  #: column type for chunk structures
    NE = "ne"  #: column type for named entities
    SRL = "srl"  #: column type for semantic role labels
    IGNORE = "ignore"  #: column type for column that should be ignored

    #: A list of all column types supported by the conll corpus reader.
    COLUMN_TYPES = (WORDS, POS, TREE, CHUNK, NE, SRL, IGNORE)

    # /////////////////////////////////////////////////////////////////
    # Constructor
    # /////////////////////////////////////////////////////////////////

    def __init__(
        self,
        root,
        fileids,
        columntypes,
        chunk_types=None,
        root_label="S",
        pos_in_tree=False,
        srl_includes_roleset=True,
        encoding="utf8",
        tree_class=Tree,
        tagset=None,
        separator=None,
    ):
        for columntype in columntypes:
            if columntype not in self.COLUMN_TYPES:
                raise ValueError("Bad column type %r" % columntype)
        if isinstance(chunk_types, str):
            chunk_types = [chunk_types]
        self._chunk_types = chunk_types
        self._colmap = {c: i for (i, c) in enumerate(columntypes)}
        self._pos_in_tree = pos_in_tree
        self._root_label = root_label  # for chunks
        self._srl_includes_roleset = srl_includes_roleset
        self._tree_class = tree_class
        CorpusReader.__init__(self, root, fileids, encoding)
        self._tagset = tagset
        self.sep = separator

    # /////////////////////////////////////////////////////////////////
    # Data Access Methods
    # /////////////////////////////////////////////////////////////////

    def words(self, fileids=None):
        self._require(self.WORDS)
        return LazyConcatenation(LazyMap(self._get_words, self._grids(fileids)))

    def sents(self, fileids=None):
        self._require(self.WORDS)
        return LazyMap(self._get_words, self._grids(fileids))

    def tagged_words(self, fileids=None, tagset=None):
        self._require(self.WORDS, self.POS)

        def get_tagged_words(grid):
            return self._get_tagged_words(grid, tagset)

        return LazyConcatenation(LazyMap(get_tagged_words, self._grids(fileids)))

    def tagged_sents(self, fileids=None, tagset=None):
        self._require(self.WORDS, self.POS)

        def get_tagged_words(grid):
            return self._get_tagged_words(grid, tagset)

        return LazyMap(get_tagged_words, self._grids(fileids))

    def chunked_words(self, fileids=None, chunk_types=None, tagset=None):
        self._require(self.WORDS, self.POS, self.CHUNK)
        if chunk_types is None:
            chunk_types = self._chunk_types

        def get_chunked_words(grid):  # capture chunk_types as local var
            return self._get_chunked_words(grid, chunk_types, tagset)

        return LazyConcatenation(LazyMap(get_chunked_words, self._grids(fileids)))

    def chunked_sents(self, fileids=None, chunk_types=None, tagset=None):
        self._require(self.WORDS, self.POS, self.CHUNK)
        if chunk_types is None:
            chunk_types = self._chunk_types

        def get_chunked_words(grid):  # capture chunk_types as local var
            return self._get_chunked_words(grid, chunk_types, tagset)

        return LazyMap(get_chunked_words, self._grids(fileids))

    def parsed_sents(self, fileids=None, pos_in_tree=None, tagset=None):
        self._require(self.WORDS, self.POS, self.TREE)
        if pos_in_tree is None:
            pos_in_tree = self._pos_in_tree

        def get_parsed_sent(grid):  # capture pos_in_tree as local var
            return self._get_parsed_sent(grid, pos_in_tree, tagset)

        return LazyMap(get_parsed_sent, self._grids(fileids))

    def srl_spans(self, fileids=None):
        self._require(self.SRL)
        return LazyMap(self._get_srl_spans, self._grids(fileids))

    def srl_instances(self, fileids=None, pos_in_tree=None, flatten=True):
        self._require(self.WORDS, self.POS, self.TREE, self.SRL)
        if pos_in_tree is None:
            pos_in_tree = self._pos_in_tree

        def get_srl_instances(grid):  # capture pos_in_tree as local var
            return self._get_srl_instances(grid, pos_in_tree)

        result = LazyMap(get_srl_instances, self._grids(fileids))
        if flatten:
            result = LazyConcatenation(result)
        return result

    def iob_words(self, fileids=None, tagset=None):
        """
        :return: a list of word/tag/IOB tuples
        :rtype: list(tuple)
        :param fileids: the list of fileids that make up this corpus
        :type fileids: None or str or list
        """
        self._require(self.WORDS, self.POS, self.CHUNK)

        def get_iob_words(grid):
            return self._get_iob_words(grid, tagset)

        return LazyConcatenation(LazyMap(get_iob_words, self._grids(fileids)))

    def iob_sents(self, fileids=None, tagset=None):
        """
        :return: a list of lists of word/tag/IOB tuples
        :rtype: list(list)
        :param fileids: the list of fileids that make up this corpus
        :type fileids: None or str or list
        """
        self._require(self.WORDS, self.POS, self.CHUNK)

        def get_iob_words(grid):
            return self._get_iob_words(grid, tagset)

        return LazyMap(get_iob_words, self._grids(fileids))

    # /////////////////////////////////////////////////////////////////
    # Grid Reading
    # /////////////////////////////////////////////////////////////////

    def _grids(self, fileids=None):
        # n.b.: we could cache the object returned here (keyed on
        # fileids), which would let us reuse the same corpus view for
        # different things (eg srl and parse trees).
        return concat(
            [
                StreamBackedCorpusView(fileid, self._read_grid_block, encoding=enc)
                for (fileid, enc) in self.abspaths(fileids, True)
            ]
        )

    def _read_grid_block(self, stream):
        grids = []
        for block in read_blankline_block(stream):
            block = block.strip()
            if not block:
                continue

            grid = [line.split(self.sep) for line in block.split("\n")]

            # If there's a docstart row, then discard. ([xx] eventually it
            # would be good to actually use it)
            if grid[0][self._colmap.get("words", 0)] == "-DOCSTART-":
                del grid[0]

            # Check that the grid is consistent.
            for row in grid:
                if len(row) != len(grid[0]):
                    raise ValueError("Inconsistent number of columns:\n%s" % block)
            grids.append(grid)
        return grids

    # /////////////////////////////////////////////////////////////////
    # Transforms
    # /////////////////////////////////////////////////////////////////
    # given a grid, transform it into some representation (e.g.,
    # a list of words or a parse tree).

    def _get_words(self, grid):
        return self._get_column(grid, self._colmap["words"])

    def _get_tagged_words(self, grid, tagset=None):
        pos_tags = self._get_column(grid, self._colmap["pos"])
        if tagset and tagset != self._tagset:
            pos_tags = [map_tag(self._tagset, tagset, t) for t in pos_tags]
        return list(zip(self._get_column(grid, self._colmap["words"]), pos_tags))

    def _get_iob_words(self, grid, tagset=None):
        pos_tags = self._get_column(grid, self._colmap["pos"])
        if tagset and tagset != self._tagset:
            pos_tags = [map_tag(self._tagset, tagset, t) for t in pos_tags]
        return list(
            zip(
                self._get_column(grid, self._colmap["words"]),
                pos_tags,
                self._get_column(grid, self._colmap["chunk"]),
            )
        )

    def _get_chunked_words(self, grid, chunk_types, tagset=None):
        # n.b.: this method is very similar to conllstr2tree.
        words = self._get_column(grid, self._colmap["words"])
        pos_tags = self._get_column(grid, self._colmap["pos"])
        if tagset and tagset != self._tagset:
            pos_tags = [map_tag(self._tagset, tagset, t) for t in pos_tags]
        chunk_tags = self._get_column(grid, self._colmap["chunk"])

        stack = [Tree(self._root_label, [])]

        for (word, pos_tag, chunk_tag) in zip(words, pos_tags, chunk_tags):
            if chunk_tag == "O":
                state, chunk_type = "O", ""
            else:
                (state, chunk_type) = chunk_tag.split("-")
            # If it's a chunk we don't care about, treat it as O.
            if chunk_types is not None and chunk_type not in chunk_types:
                state = "O"
            # Treat a mismatching I like a B.
            if state == "I" and chunk_type != stack[-1].label():
                state = "B"
            # For B or I: close any open chunks
            if state in "BO" and len(stack) == 2:
                stack.pop()
            # For B: start a new chunk.
            if state == "B":
                new_chunk = Tree(chunk_type, [])
                stack[-1].append(new_chunk)
                stack.append(new_chunk)
            # Add the word token.
            stack[-1].append((word, pos_tag))

        return stack[0]

    def _get_parsed_sent(self, grid, pos_in_tree, tagset=None):
        words = self._get_column(grid, self._colmap["words"])
        pos_tags = self._get_column(grid, self._colmap["pos"])
        if tagset and tagset != self._tagset:
            pos_tags = [map_tag(self._tagset, tagset, t) for t in pos_tags]
        parse_tags = self._get_column(grid, self._colmap["tree"])

        treestr = ""
        for (word, pos_tag, parse_tag) in zip(words, pos_tags, parse_tags):
            if word == "(":
                word = "-LRB-"
            if word == ")":
                word = "-RRB-"
            if pos_tag == "(":
                pos_tag = "-LRB-"
            if pos_tag == ")":
                pos_tag = "-RRB-"
            (left, right) = parse_tag.split("*")
            right = right.count(")") * ")"  # only keep ')'.
            treestr += f"{left} ({pos_tag} {word}) {right}"
        try:
            tree = self._tree_class.fromstring(treestr)
        except (ValueError, IndexError):
            tree = self._tree_class.fromstring(f"({self._root_label} {treestr})")

        if not pos_in_tree:
            for subtree in tree.subtrees():
                for i, child in enumerate(subtree):
                    if (
                        isinstance(child, Tree)
                        and len(child) == 1
                        and isinstance(child[0], str)
                    ):
                        subtree[i] = (child[0], child.label())

        return tree

    def _get_srl_spans(self, grid):
        """
        list of list of (start, end), tag) tuples
        """
        if self._srl_includes_roleset:
            predicates = self._get_column(grid, self._colmap["srl"] + 1)
            start_col = self._colmap["srl"] + 2
        else:
            predicates = self._get_column(grid, self._colmap["srl"])
            start_col = self._colmap["srl"] + 1

        # Count how many predicates there are.  This tells us how many
        # columns to expect for SRL data.
        num_preds = len([p for p in predicates if p != "-"])

        spanlists = []
        for i in range(num_preds):
            col = self._get_column(grid, start_col + i)
            spanlist = []
            stack = []
            for wordnum, srl_tag in enumerate(col):
                (left, right) = srl_tag.split("*")
                for tag in left.split("("):
                    if tag:
                        stack.append((tag, wordnum))
                for i in range(right.count(")")):
                    (tag, start) = stack.pop()
                    spanlist.append(((start, wordnum + 1), tag))
            spanlists.append(spanlist)

        return spanlists

    def _get_srl_instances(self, grid, pos_in_tree):
        tree = self._get_parsed_sent(grid, pos_in_tree)
        spanlists = self._get_srl_spans(grid)
        if self._srl_includes_roleset:
            predicates = self._get_column(grid, self._colmap["srl"] + 1)
            rolesets = self._get_column(grid, self._colmap["srl"])
        else:
            predicates = self._get_column(grid, self._colmap["srl"])
            rolesets = [None] * len(predicates)

        instances = ConllSRLInstanceList(tree)
        for wordnum, predicate in enumerate(predicates):
            if predicate == "-":
                continue
            # Decide which spanlist to use.  Don't assume that they're
            # sorted in the same order as the predicates (even though
            # they usually are).
            for spanlist in spanlists:
                for (start, end), tag in spanlist:
                    if wordnum in range(start, end) and tag in ("V", "C-V"):
                        break
                else:
                    continue
                break
            else:
                raise ValueError("No srl column found for %r" % predicate)
            instances.append(
                ConllSRLInstance(tree, wordnum, predicate, rolesets[wordnum], spanlist)
            )

        return instances

    # /////////////////////////////////////////////////////////////////
    # Helper Methods
    # /////////////////////////////////////////////////////////////////

    def _require(self, *columntypes):
        for columntype in columntypes:
            if columntype not in self._colmap:
                raise ValueError(
                    "This corpus does not contain a %s " "column." % columntype
                )

    @staticmethod
    def _get_column(grid, column_index):
        return [grid[i][column_index] for i in range(len(grid))]


class ConllSRLInstance:
    """
    An SRL instance from a CoNLL corpus, which identifies and
    providing labels for the arguments of a single verb.
    """

    # [xx] add inst.core_arguments, inst.argm_arguments?

    def __init__(self, tree, verb_head, verb_stem, roleset, tagged_spans):
        self.verb = []
        """A list of the word indices of the words that compose the
           verb whose arguments are identified by this instance.
           This will contain multiple word indices when multi-word
           verbs are used (e.g. 'turn on')."""

        self.verb_head = verb_head
        """The word index of the head word of the verb whose arguments
           are identified by this instance.  E.g., for a sentence that
           uses the verb 'turn on,' ``verb_head`` will be the word index
           of the word 'turn'."""

        self.verb_stem = verb_stem

        self.roleset = roleset

        self.arguments = []
        """A list of ``(argspan, argid)`` tuples, specifying the location
           and type for each of the arguments identified by this
           instance.  ``argspan`` is a tuple ``start, end``, indicating
           that the argument consists of the ``words[start:end]``."""

        self.tagged_spans = tagged_spans
        """A list of ``(span, id)`` tuples, specifying the location and
           type for each of the arguments, as well as the verb pieces,
           that make up this instance."""

        self.tree = tree
        """The parse tree for the sentence containing this instance."""

        self.words = tree.leaves()
        """A list of the words in the sentence containing this
           instance."""

        # Fill in the self.verb and self.arguments values.
        for (start, end), tag in tagged_spans:
            if tag in ("V", "C-V"):
                self.verb += list(range(start, end))
            else:
                self.arguments.append(((start, end), tag))

    def __repr__(self):
        # Originally, its:
        ##plural = 's' if len(self.arguments) != 1 else ''
        plural = "s" if len(self.arguments) != 1 else ""
        return "<ConllSRLInstance for %r with %d argument%s>" % (
            (self.verb_stem, len(self.arguments), plural)
        )

    def pprint(self):
        verbstr = " ".join(self.words[i][0] for i in self.verb)
        hdr = f"SRL for {verbstr!r} (stem={self.verb_stem!r}):\n"
        s = ""
        for i, word in enumerate(self.words):
            if isinstance(word, tuple):
                word = word[0]
            for (start, end), argid in self.arguments:
                if i == start:
                    s += "[%s " % argid
                if i == end:
                    s += "] "
            if i in self.verb:
                word = "<<%s>>" % word
            s += word + " "
        return hdr + textwrap.fill(
            s.replace(" ]", "]"), initial_indent="    ", subsequent_indent="    "
        )


class ConllSRLInstanceList(list):
    """
    Set of instances for a single sentence
    """

    def __init__(self, tree, instances=()):
        self.tree = tree
        list.__init__(self, instances)

    def __str__(self):
        return self.pprint()

    def pprint(self, include_tree=False):
        # Sanity check: trees should be the same
        for inst in self:
            if inst.tree != self.tree:
                raise ValueError("Tree mismatch!")

        # If desired, add trees:
        if include_tree:
            words = self.tree.leaves()
            pos = [None] * len(words)
            synt = ["*"] * len(words)
            self._tree2conll(self.tree, 0, words, pos, synt)

        s = ""
        for i in range(len(words)):
            # optional tree columns
            if include_tree:
                s += "%-20s " % words[i]
                s += "%-8s " % pos[i]
                s += "%15s*%-8s " % tuple(synt[i].split("*"))

            # verb head column
            for inst in self:
                if i == inst.verb_head:
                    s += "%-20s " % inst.verb_stem
                    break
            else:
                s += "%-20s " % "-"
            # Remaining columns: self
            for inst in self:
                argstr = "*"
                for (start, end), argid in inst.tagged_spans:
                    if i == start:
                        argstr = f"({argid}{argstr}"
                    if i == (end - 1):
                        argstr += ")"
                s += "%-12s " % argstr
            s += "\n"
        return s

    def _tree2conll(self, tree, wordnum, words, pos, synt):
        assert isinstance(tree, Tree)
        if len(tree) == 1 and isinstance(tree[0], str):
            pos[wordnum] = tree.label()
            assert words[wordnum] == tree[0]
            return wordnum + 1
        elif len(tree) == 1 and isinstance(tree[0], tuple):
            assert len(tree[0]) == 2
            pos[wordnum], pos[wordnum] = tree[0]
            return wordnum + 1
        else:
            synt[wordnum] = f"({tree.label()}{synt[wordnum]}"
            for child in tree:
                wordnum = self._tree2conll(child, wordnum, words, pos, synt)
            synt[wordnum - 1] += ")"
            return wordnum


class ConllChunkCorpusReader(ConllCorpusReader):
    """
    A ConllCorpusReader whose data file contains three columns: words,
    pos, and chunk.
    """

    def __init__(
        self, root, fileids, chunk_types, encoding="utf8", tagset=None, separator=None
    ):
        ConllCorpusReader.__init__(
            self,
            root,
            fileids,
            ("words", "pos", "chunk"),
            chunk_types=chunk_types,
            encoding=encoding,
            tagset=tagset,
            separator=separator,
        )