spaCy/spacy/syntax/parser.pyx

"""
MALT-style dependency parser
"""
from __future__ import unicode_literals
cimport cython
from libc.stdint cimport uint32_t, uint64_t
import random
import os.path
from os import path
import shutil
import json

from cymem.cymem cimport Pool, Address
from murmurhash.mrmr cimport hash64
from thinc.typedefs cimport weight_t, class_t, feat_t, atom_t


from util import Config

from thinc.features cimport Extractor
from thinc.features cimport Feature
from thinc.features cimport count_feats

from thinc.learner cimport LinearModel

from thinc.search cimport Beam
from thinc.search cimport MaxViolation

from ..tokens cimport Tokens, TokenC
from ..strings cimport StringStore

from .arc_eager cimport TransitionSystem, Transition
from .transition_system import OracleError

from ._state cimport State, new_state, copy_state, is_final, push_stack
from ..gold cimport GoldParse

from . import _parse_features
from ._parse_features cimport fill_context, CONTEXT_SIZE


DEBUG = False
def set_debug(val):
    global DEBUG
    DEBUG = val


cdef unicode print_state(State* s, list words):
    words = list(words) + ['EOL']
    top = words[s.stack[0]] + '_%d' % s.sent[s.stack[0]].head
    second = words[s.stack[-1]] + '_%d' % s.sent[s.stack[-1]].head
    third = words[s.stack[-2]] + '_%d' % s.sent[s.stack[-2]].head
    n0 = words[s.i] if s.i < len(words) else 'EOL'
    n1 = words[s.i + 1] if s.i+1 < len(words) else 'EOL'
    if s.ents_len:
        ent = '%s %d-%d' % (s.ent.label, s.ent.start, s.ent.end)
    else:
        ent = '-'
    return ' '.join((ent, str(s.stack_len), third, second, top, '|', n0, n1))


def get_templates(name):
    pf = _parse_features
    if name == 'ner':
        return pf.ner
    elif name == 'debug':
        return pf.unigrams
    else:
        return (pf.unigrams + pf.s0_n0 + pf.s1_n0 + pf.s0_n1 + pf.n0_n1 + \
                pf.tree_shape + pf.trigrams)


cdef class Parser:
    def __init__(self, StringStore strings, model_dir, transition_system):
        assert os.path.exists(model_dir) and os.path.isdir(model_dir)
        self.cfg = Config.read(model_dir, 'config')
        self.moves = transition_system(strings, self.cfg.labels)
        templates = get_templates(self.cfg.features)
        self.model = Model(self.moves.n_moves, templates, model_dir)

    def __call__(self, Tokens tokens):
        if tokens.length == 0:
            return 0
        if self.cfg.beam_width == 1:
            self._greedy_parse(tokens)
        else:
            self._beam_parse(tokens)

    def train(self, Tokens tokens, GoldParse gold):
        self.moves.preprocess_gold(gold)
        if self.cfg.beam_width == 1:
            return self._greedy_train(tokens, gold)
        else:
            return self._beam_train(tokens, gold)

    cdef int _greedy_parse(self, Tokens tokens) except -1:
        cdef atom_t[CONTEXT_SIZE] context
        cdef int n_feats
        cdef Pool mem = Pool()
        cdef State* state = new_state(mem, tokens.data, tokens.length)
        self.moves.initialize_state(state)
        cdef Transition guess
        while not is_final(state):
            fill_context(context, state)
            scores = self.model.score(context)
            guess = self.moves.best_valid(scores, state)
            guess.do(&guess, state)
        self.moves.finalize_state(state)
        tokens.set_parse(state.sent)

    cdef int _beam_parse(self, Tokens tokens) except -1:
        cdef Beam beam = Beam(self.moves.n_moves, self.cfg.beam_width)
        beam.initialize(_init_state, tokens.length, tokens.data)
        while not beam.is_done:
            self._advance_beam(beam, None, False)
        state = <State*>beam.at(0)
        self.moves.finalize_state(state)
        tokens.set_parse(state.sent)

    def _greedy_train(self, Tokens tokens, GoldParse gold):
        cdef Pool mem = Pool()
        cdef State* state = new_state(mem, tokens.data, tokens.length)
        self.moves.initialize_state(state)

        cdef int cost
        cdef const Feature* feats
        cdef const weight_t* scores
        cdef Transition guess
        cdef Transition best
        cdef atom_t[CONTEXT_SIZE] context
        loss = 0
        while not is_final(state):
            fill_context(context, state)
            scores = self.model.score(context)
            guess = self.moves.best_valid(scores, state)
            best = self.moves.best_gold(scores, state, gold)
            cost = guess.get_cost(&guess, state, gold)
            self.model.update(context, guess.clas, best.clas, cost)
            guess.do(&guess, state)
            loss += cost
        return loss

    def _beam_train(self, Tokens tokens, GoldParse gold_parse):
        cdef Beam pred = Beam(self.moves.n_moves, self.cfg.beam_width)
        pred.initialize(_init_state, tokens.length, tokens.data)
        cdef Beam gold = Beam(self.moves.n_moves, self.cfg.beam_width)
        gold.initialize(_init_state, tokens.length, tokens.data)

        violn = MaxViolation()
        while not pred.is_done and not gold.is_done:
            self._advance_beam(pred, gold_parse, False)
            self._advance_beam(gold, gold_parse, True)
            violn.check(pred, gold)
        counts = {}
        if pred._states[0].loss >= 1:
            self._count_feats(counts, tokens, violn.g_hist, 1)
            self._count_feats(counts, tokens, violn.p_hist, -1)
        self.model._model.update(counts)
        return pred._states[0].loss

    def _advance_beam(self, Beam beam, GoldParse gold, bint follow_gold):
        cdef atom_t[CONTEXT_SIZE] context
        cdef State* state
        cdef int i, j, cost
        cdef bint is_valid
        cdef const Transition* move
        for i in range(beam.size):
            state = <State*>beam.at(i)
            fill_context(context, state)
            scores = self.model.score(context)
            validities = self.moves.get_valid(state)
            if gold is None:
                for j in range(self.moves.n_moves):
                    beam.set_cell(i, j, scores[j], validities[j], 0)
            elif not follow_gold:
                for j in range(self.moves.n_moves):
                    move = &self.moves.c[j]
                    cost = move.get_cost(move, state, gold)
                    beam.set_cell(i, j, scores[j], validities[j], cost)
            else:
                for j in range(self.moves.n_moves):
                    move = &self.moves.c[j]
                    cost = move.get_cost(move, state, gold)
                    beam.set_cell(i, j, scores[j], cost == 0, cost)
        beam.advance(_transition_state, <void*>self.moves.c)
        beam.check_done(_check_final_state, NULL)

    def _count_feats(self, dict counts, Tokens tokens, list hist, int inc):
        cdef atom_t[CONTEXT_SIZE] context
        cdef Pool mem = Pool()
        cdef State* state = new_state(mem, tokens.data, tokens.length)
        self.moves.initialize_state(state)

        cdef class_t clas
        cdef int n_feats
        for clas in hist:
            if is_final(state):
                break
            fill_context(context, state)
            feats = self.model._extractor.get_feats(context, &n_feats)
            count_feats(counts.setdefault(clas, {}), feats, n_feats, inc)
            self.moves.c[clas].do(&self.moves.c[clas], state)


# These are passed as callbacks to thinc.search.Beam

cdef int _transition_state(void* _dest, void* _src, class_t clas, void* _moves) except -1:
    dest = <State*>_dest
    src = <const State*>_src
    moves = <const Transition*>_moves
    copy_state(dest, src)
    moves[clas].do(&moves[clas], dest)


cdef void* _init_state(Pool mem, int length, void* tokens) except NULL:
    state = new_state(mem, <const TokenC*>tokens, length)
    push_stack(state)
    return state


cdef int _check_final_state(void* state, void* extra_args) except -1:
    return is_final(<State*>state)