* Fix import of sense tagger

.gitignore

@@ -17,6 +17,8 @@ models/
 spacy/syntax/*.cpp
 spacy/syntax/*.html
 spacy/en/*.cpp
+spacy/wsd/*.cpp
+spacy/wsd/*.html
 spacy/en/data/*
 spacy/*.cpp
 spacy/ner/*.cpp

bin/init_model.py

@@ -22,15 +22,17 @@ from shutil import copyfile
 from shutil import copytree
 import codecs
 from collections import defaultdict
+import json
 
 from spacy.en import get_lex_props
 from spacy.en.lemmatizer import Lemmatizer
 from spacy.vocab import Vocab
 from spacy.vocab import write_binary_vectors
 
-from spacy.parts_of_speech import NOUN, VERB, ADJ
+from spacy.parts_of_speech import NOUN, VERB, ADJ, ADV
 
 import spacy.senses
+from spacy.munge import read_wordnet
 
 
 def setup_tokenizer(lang_data_dir, tok_dir):
@@ -80,17 +82,13 @@ def _read_probs(loc):
 
 def _read_senses(loc):
     lexicon = defaultdict(lambda: defaultdict(list))
-    sense_names = dict((s, i) for i, s in enumerate(spacy.senses.STRINGS))
-    pos_ids = {'noun': NOUN, 'verb': VERB, 'adjective': ADJ}
+    pos_tags = [None, NOUN, VERB, ADJ, None, None]
     for line in codecs.open(str(loc), 'r', 'utf8'):
-        sense_strings = line.split()
-        word = sense_strings.pop(0)
-        for sense in sense_strings:
-            pos, sense = sense[3:].split('.')
-            sense_name = '%s_%s' % (pos[0].upper(), sense.lower())
-            if sense_name != 'N_tops':
-                sense_id = sense_names[sense_name]
-                lexicon[word][pos_ids[pos]].append(sense_id)
+        sense_key, synset_offset, sense_number, tag_cnt = line.split()
+        lemma, lex_sense = sense_key.split('%')
+        ss_type, lex_filenum, lex_id, head_word, head_id = lex_sense.split(':')
+        pos = pos_tags[int(ss_type)]
+        lexicon[lemma][pos].append(int(lex_filenum) + 1)
     return lexicon
 
 
@@ -103,7 +101,7 @@ def setup_vocab(src_dir, dst_dir):
         write_binary_vectors(str(vectors_src), str(dst_dir / 'vec.bin'))
     vocab = Vocab(data_dir=None, get_lex_props=get_lex_props)
     clusters = _read_clusters(src_dir / 'clusters.txt')
-    senses = _read_senses(src_dir / 'supersenses.txt')
+    senses = _read_senses(src_dir / 'wordnet' / 'index.sense')
     probs = _read_probs(src_dir / 'words.sgt.prob')
     for word in set(clusters).union(set(senses)):
         if word not in probs:
@@ -112,28 +110,24 @@ def setup_vocab(src_dir, dst_dir):
     lexicon = []
     for word, prob in reversed(sorted(probs.items(), key=lambda item: item[1])):
         entry = get_lex_props(word)
-        if word in clusters or float(prob) >= -17:
+        if word in clusters or word in senses or float(prob) >= -17:
             entry['prob'] = float(prob)
             cluster = clusters.get(word, '0')
             # Decode as a little-endian string, so that we can do & 15 to get
             # the first 4 bits. See _parse_features.pyx
             entry['cluster'] = int(cluster[::-1], 2)
             orth_senses = set()
-            lemmas = []
+            orth_senses.update(senses[word.lower()][None])
             for pos in [NOUN, VERB, ADJ]:
                 for lemma in lemmatizer(word.lower(), pos):
-                    lemmas.append(lemma)
                     orth_senses.update(senses[lemma][pos])
-            if word.lower() == 'dogging':
-                print word
-                print lemmas
-                print [spacy.senses.STRINGS[si] for si in orth_senses]
             entry['senses'] = list(sorted(orth_senses))
             vocab[word] = entry
     vocab.dump(str(dst_dir / 'lexemes.bin'))
     vocab.strings.dump(str(dst_dir / 'strings.txt'))
 
 
+
 def main(lang_data_dir, corpora_dir, model_dir):
     model_dir = Path(model_dir)
     lang_data_dir = Path(lang_data_dir)
@@ -147,8 +141,12 @@ def main(lang_data_dir, corpora_dir, model_dir):
 
     setup_tokenizer(lang_data_dir, model_dir / 'tokenizer')
     setup_vocab(corpora_dir, model_dir / 'vocab')
+
     if not (model_dir / 'wordnet').exists():
         copytree(str(corpora_dir / 'wordnet'), str(model_dir / 'wordnet'))
+    ss_probs = read_wordnet.make_supersense_dict(str(corpora_dir / 'wordnet'))
+    with codecs.open(str(model_dir / 'wordnet' / 'supersenses.json'), 'w', 'utf8') as file_:
+        json.dump(ss_probs, file_)
 
 
 if __name__ == '__main__':

bin/parser/nn_train.py

@@ -0,0 +1,261 @@
+#!/usr/bin/env python
+from __future__ import division
+from __future__ import unicode_literals
+
+import os
+from os import path
+import shutil
+import codecs
+import random
+
+import plac
+import cProfile
+import pstats
+import re
+
+import spacy.util
+from spacy.en import English
+from spacy.en.pos import POS_TEMPLATES, POS_TAGS, setup_model_dir
+
+from spacy.syntax.util import Config
+from spacy.gold import read_json_file
+from spacy.gold import GoldParse
+
+from spacy.scorer import Scorer
+
+from spacy.syntax.parser import Parser, get_templates
+from spacy._theano import TheanoModel
+
+import theano
+import theano.tensor as T
+
+from theano.printing import Print
+
+import numpy
+from collections import OrderedDict, defaultdict
+
+
+theano.config.profile = False
+theano.config.floatX = 'float32'
+floatX = theano.config.floatX
+
+
+def L1(L1_reg, *weights):
+    return L1_reg * sum(abs(w).sum() for w in weights)
+
+
+def L2(L2_reg, *weights):
+    return L2_reg * sum((w ** 2).sum() for w in weights)
+
+
+def rms_prop(loss, params, eta=1.0, rho=0.9, eps=1e-6):
+    updates = OrderedDict()
+    for param in params:
+        value = param.get_value(borrow=True)
+        accu = theano.shared(np.zeros(value.shape, dtype=value.dtype),
+                             broadcastable=param.broadcastable)
+
+        grad = T.grad(loss, param)
+        accu_new = rho * accu + (1 - rho) * grad ** 2
+        updates[accu] = accu_new
+        updates[param] = param - (eta * grad / T.sqrt(accu_new + eps))
+    return updates
+
+
+def relu(x):
+    return x * (x > 0)
+
+
+def feed_layer(activation, weights, bias, input_):
+    return activation(T.dot(input_, weights) + bias)
+
+
+def init_weights(n_in, n_out):
+    rng = numpy.random.RandomState(1235)
+    
+    weights = numpy.asarray(
+        rng.standard_normal(size=(n_in, n_out)) * numpy.sqrt(2.0 / n_in),
+        dtype=theano.config.floatX
+    )
+    bias = numpy.zeros((n_out,), dtype=theano.config.floatX)
+    return [wrapper(weights, name='W'), wrapper(bias, name='b')]
+
+
+def compile_model(n_classes, n_hidden, n_in, optimizer):
+    x = T.vector('x') 
+    costs = T.ivector('costs')
+    loss = T.scalar('loss')
+
+    maxent_W, maxent_b = init_weights(n_hidden, n_classes)
+    hidden_W, hidden_b = init_weights(n_in, n_hidden)
+
+    # Feed the inputs forward through the network
+    p_y_given_x = feed_layer(
+                    T.nnet.softmax,
+                    maxent_W,
+                    maxent_b,
+                      feed_layer(
+                        relu,
+                        hidden_W,
+                        hidden_b,
+                        x))
+
+    loss = -T.log(T.sum(p_y_given_x[0] * T.eq(costs, 0)) + 1e-8)
+
+    train_model = theano.function(
+        name='train_model',
+        inputs=[x, costs],
+        outputs=[p_y_given_x[0], T.grad(loss, x), loss],
+        updates=optimizer(loss, [maxent_W, maxent_b, hidden_W, hidden_b]),
+        on_unused_input='warn'
+    )
+
+    evaluate_model = theano.function(
+        name='evaluate_model',
+        inputs=[x],
+        outputs=[
+            feed_layer(
+              T.nnet.softmax,
+              maxent_W,
+              maxent_b,
+              feed_layer(
+                relu,
+                hidden_W,
+                hidden_b,
+                x
+              )
+            )[0]
+        ]
+    )
+    return train_model, evaluate_model
+
+
+def score_model(scorer, nlp, annot_tuples, verbose=False):
+    tokens = nlp.tokenizer.tokens_from_list(annot_tuples[1])
+    nlp.tagger(tokens)
+    nlp.parser(tokens)
+    gold = GoldParse(tokens, annot_tuples)
+    scorer.score(tokens, gold, verbose=verbose)
+
+
+def train(Language, gold_tuples, model_dir, n_iter=15, feat_set=u'basic',
+          eta=0.01, mu=0.9, nv_hidden=100, nv_word=10, nv_tag=10, nv_label=10,
+          seed=0, n_sents=0,  verbose=False):
+
+    dep_model_dir = path.join(model_dir, 'deps')
+    pos_model_dir = path.join(model_dir, 'pos')
+    if path.exists(dep_model_dir):
+        shutil.rmtree(dep_model_dir)
+    if path.exists(pos_model_dir):
+        shutil.rmtree(pos_model_dir)
+    os.mkdir(dep_model_dir)
+    os.mkdir(pos_model_dir)
+    setup_model_dir(sorted(POS_TAGS.keys()), POS_TAGS, POS_TEMPLATES, pos_model_dir)
+
+    Config.write(dep_model_dir, 'config',
+        seed=seed,
+        templates=tuple(),
+        labels=Language.ParserTransitionSystem.get_labels(gold_tuples),
+        vector_lengths=(nv_word, nv_tag, nv_label),
+        hidden_nodes=nv_hidden,
+        eta=eta,
+        mu=mu
+    )
+  
+    # Bake-in hyper-parameters
+    optimizer = lambda loss, params: rms_prop(loss, params, eta=eta, rho=rho, eps=eps)
+    nlp = Language(data_dir=model_dir)
+    n_classes = nlp.parser.model.n_classes
+    train, predict = compile_model(n_classes, nv_hidden, n_in, optimizer)
+    nlp.parser.model = TheanoModel(n_classes, input_spec, train,
+                                   predict, model_loc)
+ 
+    if n_sents > 0:
+        gold_tuples = gold_tuples[:n_sents]
+    print "Itn.\tP.Loss\tUAS\tTag %\tToken %"
+    log_loc = path.join(model_dir, 'job.log')
+    for itn in range(n_iter):
+        scorer = Scorer()
+        loss = 0
+        for _, sents in gold_tuples:
+            for annot_tuples, ctnt in sents:
+                if len(annot_tuples[1]) == 1:
+                    continue
+                score_model(scorer, nlp, annot_tuples)
+                tokens = nlp.tokenizer.tokens_from_list(annot_tuples[1])
+                nlp.tagger(tokens)
+                gold = GoldParse(tokens, annot_tuples, make_projective=True)
+                assert gold.is_projective
+                loss += nlp.parser.train(tokens, gold)
+                nlp.tagger.train(tokens, gold.tags)
+        random.shuffle(gold_tuples)
+        logline = '%d:\t%d\t%.3f\t%.3f\t%.3f' % (itn, loss, scorer.uas,
+                                                 scorer.tags_acc,
+                                                 scorer.token_acc)
+        print logline
+        with open(log_loc, 'aw') as file_:
+            file_.write(logline + '\n')
+    nlp.parser.model.end_training()
+    nlp.tagger.model.end_training()
+    nlp.vocab.strings.dump(path.join(model_dir, 'vocab', 'strings.txt'))
+    return nlp
+
+
+def evaluate(nlp, gold_tuples, gold_preproc=True):
+    scorer = Scorer()
+    for raw_text, sents in gold_tuples:
+        for annot_tuples, brackets in sents:
+            tokens = nlp.tokenizer.tokens_from_list(annot_tuples[1])
+            nlp.tagger(tokens)
+            nlp.parser(tokens)
+            gold = GoldParse(tokens, annot_tuples)
+            scorer.score(tokens, gold)
+    return scorer
+
+
+@plac.annotations(
+    train_loc=("Location of training file or directory"),
+    dev_loc=("Location of development file or directory"),
+    model_dir=("Location of output model directory",),
+    eval_only=("Skip training, and only evaluate", "flag", "e", bool),
+    n_sents=("Number of training sentences", "option", "n", int),
+    n_iter=("Number of training iterations", "option", "i", int),
+    verbose=("Verbose error reporting", "flag", "v", bool),
+
+    nv_word=("Word vector length", "option", "W", int),
+    nv_tag=("Tag vector length", "option", "T", int),
+    nv_label=("Label vector length", "option", "L", int),
+    nv_hidden=("Hidden nodes length", "option", "H", int),
+    eta=("Learning rate", "option", "E", float),
+    mu=("Momentum", "option", "M", float),
+)
+def main(train_loc, dev_loc, model_dir, n_sents=0, n_iter=15, verbose=False,
+         nv_word=10, nv_tag=10, nv_label=10, nv_hidden=10,
+         eta=0.1, mu=0.9, eval_only=False):
+
+
+
+
+    gold_train = list(read_json_file(train_loc, lambda doc: 'wsj' in doc['id']))
+
+    nlp = train(English, gold_train, model_dir,
+               feat_set='embed',
+               eta=eta, mu=mu,
+               nv_word=nv_word, nv_tag=nv_tag, nv_label=nv_label, nv_hidden=nv_hidden,
+               n_sents=n_sents, n_iter=n_iter,
+               verbose=verbose)
+
+    scorer = evaluate(nlp, list(read_json_file(dev_loc)))
+    
+    print 'TOK', 100-scorer.token_acc
+    print 'POS', scorer.tags_acc
+    print 'UAS', scorer.uas
+    print 'LAS', scorer.las
+
+    print 'NER P', scorer.ents_p
+    print 'NER R', scorer.ents_r
+    print 'NER F', scorer.ents_f
+
+
+if __name__ == '__main__':
+    plac.call(main)

bin/parser/train.py

@@ -18,6 +18,7 @@ from spacy.en import English
 from spacy.en.pos import POS_TEMPLATES, POS_TAGS, setup_model_dir
 
 from spacy.syntax.util import Config
+from spacy.syntax.parser import get_templates
 from spacy.gold import read_json_file
 from spacy.gold import GoldParse
 
@@ -95,10 +96,10 @@ def train(Language, gold_tuples, model_dir, n_iter=15, feat_set=u'basic',
 
     setup_model_dir(sorted(POS_TAGS.keys()), POS_TAGS, POS_TEMPLATES, pos_model_dir)
 
-    Config.write(dep_model_dir, 'config', features=feat_set, seed=seed,
-                 labels=Language.ParserTransitionSystem.get_labels(gold_tuples),
+    Config.write(dep_model_dir, 'config', templates=get_templates(feat_set),
+                 seed=seed, labels=Language.ParserTransitionSystem.get_labels(gold_tuples),
                  beam_width=beam_width)
-    Config.write(ner_model_dir, 'config', features='ner', seed=seed,
+    Config.write(ner_model_dir, 'config', templates=get_templates('ner'), seed=seed,
                  labels=Language.EntityTransitionSystem.get_labels(gold_tuples),
                  beam_width=0)
 

bin/prepare_treebank.py

@@ -9,7 +9,8 @@ doc: {
             start: int,
             tag: string,
             head: int,
-            dep: string}],
+            dep: string,
+            ssenses: [int]}],
         ner: [{
             start: int,
             end: int,
@@ -33,6 +34,7 @@ from collections import defaultdict
 from spacy.munge import read_ptb
 from spacy.munge import read_conll
 from spacy.munge import read_ner
+from spacy.munge import read_wordnet
 
 
 def _iter_raw_files(raw_loc):
@@ -41,7 +43,7 @@ def _iter_raw_files(raw_loc):
         yield f
 
 
-def format_doc(file_id, raw_paras, ptb_text, dep_text, ner_text):
+def format_doc(file_id, raw_paras, ptb_text, dep_text, ner_text, senses):
     ptb_sents = read_ptb.split(ptb_text)
     dep_sents = read_conll.split(dep_text)
     if len(ptb_sents) != len(dep_sents):
@@ -54,7 +56,8 @@ def format_doc(file_id, raw_paras, ptb_text, dep_text, ner_text):
     i = 0
     doc = {'id': file_id}
     if raw_paras is None:
-        doc['paragraphs'] = [format_para(None, ptb_sents, dep_sents, ner_sents)]
+        doc['paragraphs'] = [format_para(None, ptb_sents, dep_sents, ner_sents,
+                                         [senses[j] for j in range(len(ptb_sents))])]
         #for ptb_sent, dep_sent, ner_sent in zip(ptb_sents, dep_sents, ner_sents):
         #    doc['paragraphs'].append(format_para(None, [ptb_sent], [dep_sent], [ner_sent]))
     else:
@@ -64,18 +67,21 @@ def format_doc(file_id, raw_paras, ptb_text, dep_text, ner_text):
                         ' '.join(raw_sents).replace('<SEP>', ''),
                         ptb_sents[i:i+len(raw_sents)],
                         dep_sents[i:i+len(raw_sents)],
-                        ner_sents[i:i+len(raw_sents)])
+                        ner_sents[i:i+len(raw_sents)],
+                        [senses[j] for j in range(i, i+len(raw_sents))]
+                    )
             if para['sentences']:
                 doc['paragraphs'].append(para)
             i += len(raw_sents)
     return doc
 
 
-def format_para(raw_text, ptb_sents, dep_sents, ner_sents):
+def format_para(raw_text, ptb_sents, dep_sents, ner_sents, ssenses):
     para = {'raw': raw_text, 'sentences': []}
     offset = 0
-    assert len(ptb_sents) == len(dep_sents) == len(ner_sents)
-    for ptb_text, dep_text, ner_text in zip(ptb_sents, dep_sents, ner_sents):
+
+    assert len(ptb_sents) == len(dep_sents) == len(ner_sents) == len(ssenses)
+    for ptb_text, dep_text, ner_text, sense_sent in zip(ptb_sents, dep_sents, ner_sents, ssenses):
         _, deps = read_conll.parse(dep_text, strip_bad_periods=True)
         if deps and 'VERB' in [t['tag'] for t in deps]:
             continue
@@ -87,14 +93,14 @@ def format_para(raw_text, ptb_sents, dep_sents, ner_sents):
         # Necessary because the ClearNLP converter deletes EDITED words.
         if len(ner) != len(deps):
             ner = ['-' for _ in deps]
-        para['sentences'].append(format_sentence(deps, ner, brackets))
+        para['sentences'].append(format_sentence(deps, ner, brackets, sense_sent))
     return para
 
 
-def format_sentence(deps, ner, brackets):
+def format_sentence(deps, ner, brackets, senses):
     sent = {'tokens': [], 'brackets': []}
     for token_id, (token, token_ent) in enumerate(zip(deps, ner)):
-        sent['tokens'].append(format_token(token_id, token, token_ent))
+        sent['tokens'].append(format_token(token_id, token, token_ent, senses))
 
     for label, start, end in brackets:
         if start != end:
@@ -105,16 +111,20 @@ def format_sentence(deps, ner, brackets):
     return sent
 
 
-def format_token(token_id, token, ner):
+def format_token(token_id, token, ner, senses):
     assert token_id == token['id']
     head = (token['head'] - token_id) if token['head'] != -1 else 0
+    # TODO: Sense data currently broken, due to alignment problems. Also should
+    # output OntoNotes groups, not WordNet supersenses. Don't print the information
+    # until this is fixed.
     return {
         'id': token_id,
         'orth': token['word'],
         'tag': token['tag'],
         'head': head,
         'dep': token['dep'],
-        'ner': ner}
+        'ner': ner,
+    }
 
 
 def read_file(*pieces):
@@ -132,7 +142,7 @@ def get_file_names(section_dir, subsection):
     return list(sorted(set(filenames)))
 
 
-def read_wsj_with_source(onto_dir, raw_dir):
+def read_wsj_with_source(onto_dir, raw_dir, wn_ssenses):
     # Now do WSJ, with source alignment
     onto_dir = path.join(onto_dir, 'data', 'english', 'annotations', 'nw', 'wsj')
     docs = {}
@@ -147,12 +157,14 @@ def read_wsj_with_source(onto_dir, raw_dir):
             ptb = read_file(onto_dir, section, '%s.parse' % filename)
             dep = read_file(onto_dir, section, '%s.parse.dep' % filename)
             ner = read_file(onto_dir, section, '%s.name' % filename)
-            if ptb is not None and dep is not None:
-                docs[filename] = format_doc(filename, raw_paras, ptb, dep, ner)
+            wsd = read_senses(path.join(onto_dir, section, '%s.sense' % filename), wn_ssenses)
+            if ptb is not None and dep is not None: # TODO: This is bad right?
+                wsd = [wsd[sent_id] for sent_id in range(len(ner))]
+                docs[filename] = format_doc(filename, raw_paras, ptb, dep, ner, wsd)
     return docs
 
 
-def get_doc(onto_dir, file_path, wsj_docs):
+def get_doc(onto_dir, file_path, wsj_docs, wn_ssenses):
     filename = file_path.rsplit('/', 1)[1]
     if filename in wsj_docs:
         return wsj_docs[filename]
@@ -160,8 +172,9 @@ def get_doc(onto_dir, file_path, wsj_docs):
         ptb = read_file(onto_dir, file_path + '.parse')
         dep = read_file(onto_dir, file_path + '.parse.dep')
         ner = read_file(onto_dir, file_path + '.name')
+        wsd = read_senses(file_path + '.sense', wn_ssenses)
         if ptb is not None and dep is not None:
-            return format_doc(filename, None, ptb, dep, ner)
+            return format_doc(filename, None, ptb, dep, ner, wsd)
         else:
             return None
 
@@ -170,14 +183,29 @@ def read_ids(loc):
     return open(loc).read().strip().split('\n')
 
 
-def main(onto_dir, raw_dir, out_dir):
-    wsj_docs = read_wsj_with_source(onto_dir, raw_dir)
+def read_senses(loc, og_to_ssense):
+    senses = defaultdict(lambda: defaultdict(list))
+    if not path.exists(loc):
+        return senses
+    for line in open(loc):
+        pieces = line.split()
+        sent_id = int(pieces[1])
+        tok_id = int(pieces[2])
+        lemma, pos = pieces[3].split('-')
+        group_num = int(float(pieces[-1]))
+        senses[sent_id][tok_id] = list(sorted(og_to_ssense[(lemma, pos, group_num)]))
+    return senses
+
+
+def main(wordnet_dir, onto_dir, raw_dir, out_dir):
+    wn_ssenses = read_wordnet.get_og_to_ssenses(wordnet_dir, onto_dir)
+    wsj_docs = read_wsj_with_source(onto_dir, raw_dir, wn_ssenses)
 
     for partition in ('train', 'test', 'development'):
         ids = read_ids(path.join(onto_dir, '%s.id' % partition))
         docs_by_genre = defaultdict(list)
         for file_path in ids:
-            doc = get_doc(onto_dir, file_path, wsj_docs)
+            doc = get_doc(onto_dir, file_path, wsj_docs, wn_ssenses)
             if doc is not None:
                 genre = file_path.split('/')[3]
                 docs_by_genre[genre].append(doc)

bin/wsd/unsupervised_train.py

@@ -0,0 +1,80 @@
+#!/usr/bin/env python
+from __future__ import division
+from __future__ import unicode_literals
+
+import os
+from os import path
+import random
+import shutil
+
+import plac
+
+from spacy.munge.corpus import DocsDB
+from spacy.munge.read_semcor import read_semcor
+
+from spacy.en import English
+from spacy.syntax.util import Config
+
+
+def score_model(nlp, semcor_docs):
+    n_right = 0
+    n_wrong = 0
+    n_multi = 0
+    for dnum, paras in semcor_docs:
+        for pnum, para in paras:
+            for snum, sent in para:
+                words = [t.orth for t in sent]
+                if len(words) < 2:
+                    continue
+                tokens = nlp.tokenizer.tokens_from_list(words)
+                nlp.tagger(tokens)
+                nlp.parser(tokens)
+                nlp.senser(tokens)
+                for i, token in enumerate(tokens):
+                    if '_' in sent[i].orth:
+                        n_multi += 1
+                    elif sent[i].supersense != 'NO_SENSE':
+                        n_right += token.sense_ == sent[i].supersense
+                        n_wrong += token.sense_ != sent[i].supersense
+    return n_right / (n_right + n_wrong)
+
+
+def train(Language, model_dir, train_docs, dev_docs,
+          report_every=1000, n_docs=1000, seed=0):
+    wsd_model_dir = path.join(model_dir, 'wsd')
+    if path.exists(wsd_model_dir):
+        shutil.rmtree(wsd_model_dir)
+    os.mkdir(wsd_model_dir)
+    
+    Config.write(wsd_model_dir, 'config', seed=seed)
+
+    nlp = Language(data_dir=model_dir, load_vectors=False)
+
+    loss = 0
+    n_tokens = 0
+    for i, doc in enumerate(train_docs):
+        tokens = nlp(doc, parse=True, entity=False)
+        loss += nlp.senser.train(tokens)
+        n_tokens += len(tokens)
+        if i and i % report_every == 0:
+            acc = score_model(nlp, dev_docs)
+            print i, loss / n_tokens, acc
+    nlp.senser.end_training()
+    nlp.vocab.strings.dump(path.join(model_dir, 'vocab', 'strings.txt'))
+
+
+@plac.annotations(
+    train_loc=("Location of the documents SQLite database"),
+    dev_loc=("Location of the SemCor corpus directory"),
+    model_dir=("Location of the models directory"),
+    n_docs=("Number of training documents", "option", "n", int),
+    seed=("Random seed", "option", "s", int),
+)
+def main(train_loc, dev_loc, model_dir, n_docs=1000000, seed=0):
+    train_docs = DocsDB(train_loc, limit=n_docs)
+    dev_docs = read_semcor(dev_loc)
+    train(English, model_dir, train_docs, dev_docs, report_every=100, seed=seed)
+
+
+if __name__ == '__main__':
+    plac.call(main)

docs/source/index.rst

@@ -14,7 +14,21 @@ spaCy: Industrial-strength NLP
 
 .. _Version 0.87 released: updates.html
 
-`spaCy`_ is a new library for text processing in Python and Cython.
+`spaCy`_ is a new library for text processing in Python and Cython.  It is designed
+to be production quality, general purpose, and up-to-the-minute with the latest
+research.  It offers similar functionality to Stanford's CoreNLP, but is
+faster, more accurate, and offers commercial licensing options (you can also
+use it under the AGPL).  
+
+If you're trying to do something that's never been done before, spaCy can
+instantly fast-forward you to all the best of what *has* been done --- even if
+it's only been done in a paper that was published two months ago.
+
+
+It is a production-quality implementation of underlying, general-purpose
+natural language understandi
+Its mission is to make a prod the latest natural language understanding
+research into practice, by making a production-quality implementation 
 I wrote it because I think small companies are terrible at
 natural language processing (NLP).  Or rather:
 small companies are using terrible NLP technology.

setup.py

@@ -151,13 +151,14 @@ MOD_NAMES = ['spacy.parts_of_speech', 'spacy.strings',
              'spacy.lexeme', 'spacy.vocab', 'spacy.tokens', 'spacy.spans',
              'spacy.morphology', 
              'spacy.syntax.stateclass', 
-             'spacy._ml', 'spacy.tokenizer', 'spacy.en.attrs',
+             'spacy._ml', 'spacy._theano',
+             'spacy.tokenizer', 'spacy.en.attrs',
              'spacy.en.pos', 'spacy.syntax.parser', 
              'spacy.syntax.transition_system',
              'spacy.syntax.arc_eager',
              'spacy.syntax._parse_features',
              'spacy.gold', 'spacy.orth', 
-             'spacy.senses',
+             'spacy.wsd.supersenses', 'spacy.wsd.supersense_tagger',
              'spacy.syntax.ner']
 
 

spacy/_bu_nn.pyx

@@ -0,0 +1,490 @@
+"""Feed-forward neural network, using Thenao."""
+
+import os
+import sys
+import time
+
+import numpy
+
+import theano
+import theano.tensor as T
+import gzip
+import cPickle
+
+
+def load_data(dataset):
+    ''' Loads the dataset
+
+    :type dataset: string
+    :param dataset: the path to the dataset (here MNIST)
+    '''
+
+    #############
+    # LOAD DATA #
+    #############
+
+    # Download the MNIST dataset if it is not present
+    data_dir, data_file = os.path.split(dataset)
+    if data_dir == "" and not os.path.isfile(dataset):
+        # Check if dataset is in the data directory.
+        new_path = os.path.join(
+            os.path.split(__file__)[0],
+            "..",
+            "data",
+            dataset
+        )
+        if os.path.isfile(new_path) or data_file == 'mnist.pkl.gz':
+            dataset = new_path
+
+    if (not os.path.isfile(dataset)) and data_file == 'mnist.pkl.gz':
+        import urllib
+        origin = (
+            'http://www.iro.umontreal.ca/~lisa/deep/data/mnist/mnist.pkl.gz'
+        )
+        print 'Downloading data from %s' % origin
+        urllib.urlretrieve(origin, dataset)
+
+    print '... loading data'
+
+    # Load the dataset
+    f = gzip.open(dataset, 'rb')
+    train_set, valid_set, test_set = cPickle.load(f)
+    f.close()
+    #train_set, valid_set, test_set format: tuple(input, target)
+    #input is an numpy.ndarray of 2 dimensions (a matrix), 
+    #each row corresponding to an example. target is a
+    #numpy.ndarray of 1 dimension (vector)) that have the same length as
+    #the number of rows in the input. It should give the target
+    #target to the example with the same index in the input.
+
+    def shared_dataset(data_xy, borrow=True):
+        """ Function that loads the dataset into shared variables
+
+        The reason we store our dataset in shared variables is to allow
+        Theano to copy it into the GPU memory (when code is run on GPU).
+        Since copying data into the GPU is slow, copying a minibatch everytime
+        is needed (the default behaviour if the data is not in a shared
+        variable) would lead to a large decrease in performance.
+        """
+        data_x, data_y = data_xy
+        shared_x = theano.shared(numpy.asarray(data_x, dtype=theano.config.floatX),
+                                 borrow=borrow)
+        shared_y = theano.shared(numpy.asarray(data_y, dtype=theano.config.floatX),
+                                 borrow=borrow)
+        # When storing data on the GPU it has to be stored as floats
+        # therefore we will store the labels as ``floatX`` as well
+        # (``shared_y`` does exactly that). But during our computations
+        # we need them as ints (we use labels as index, and if they are
+        # floats it doesn't make sense) therefore instead of returning
+        # ``shared_y`` we will have to cast it to int. This little hack
+        # lets ous get around this issue
+        return shared_x, T.cast(shared_y, 'int32')
+
+    test_set_x, test_set_y = shared_dataset(test_set)
+    valid_set_x, valid_set_y = shared_dataset(valid_set)
+    train_set_x, train_set_y = shared_dataset(train_set)
+
+    rval = [(train_set_x, train_set_y), (valid_set_x, valid_set_y),
+            (test_set_x, test_set_y)]
+    return rval
+
+
+class LogisticRegression(object):
+    """Multi-class Logistic Regression Class
+
+    The logistic regression is fully described by a weight matrix :math:`W`
+    and bias vector :math:`b`. Classification is done by projecting data
+    points onto a set of hyperplanes, the distance to which is used to
+    determine a class membership probability.
+    """
+
+    def __init__(self, input, n_in, n_out):
+        """ Initialize the parameters of the logistic regression
+
+        :type input: theano.tensor.TensorType
+        :param input: symbolic variable that describes the input of the
+                      architecture (one minibatch)
+
+        :type n_in: int
+        :param n_in: number of input units, the dimension of the space in
+                     which the datapoints lie
+
+        :type n_out: int
+        :param n_out: number of output units, the dimension of the space in
+                      which the labels lie
+
+        """
+        # start-snippet-1
+        # initialize with 0 the weights W as a matrix of shape (n_in, n_out)
+        self.W = theano.shared(
+            value=numpy.zeros((n_in, n_out),
+                dtype=theano.config.floatX
+            ),
+            name='W',
+            borrow=True
+        )
+        # initialize the baises b as a vector of n_out 0s
+        self.b = theano.shared(
+            value=numpy.zeros(
+                (n_out,),
+                dtype=theano.config.floatX
+            ),
+            name='b',
+            borrow=True
+        )
+
+        # symbolic expression for computing the matrix of class-membership
+        # probabilities
+        # Where:
+        # W is a matrix where column-k represent the separation hyper plain for
+        # class-k
+        # x is a matrix where row-j  represents input training sample-j
+        # b is a vector where element-k represent the free parameter of hyper
+        # plain-k
+        self.p_y_given_x = T.nnet.softmax(T.dot(input, self.W) + self.b)
+
+        # symbolic description of how to compute prediction as class whose
+        # probability is maximal
+        self.y_pred = T.argmax(self.p_y_given_x, axis=1)
+        # end-snippet-1
+
+        # parameters of the model
+        self.params = [self.W, self.b]
+
+    def neg_ll(self, y):
+        """Return the mean of the negative log-likelihood of the prediction
+        of this model under a given target distribution.
+
+        .. math::
+
+            \frac{1}{|\mathcal{D}|} \mathcal{L} (\theta=\{W,b\}, \mathcal{D}) =
+            \frac{1}{|\mathcal{D}|} \sum_{i=0}^{|\mathcal{D}|}
+                \log(P(Y=y^{(i)}|x^{(i)}, W,b)) \\
+            \ell (\theta=\{W,b\}, \mathcal{D})
+
+        :type y: theano.tensor.TensorType
+        :param y: corresponds to a vector that gives for each example the
+                  correct label
+
+        Note: we use the mean instead of the sum so that
+              the learning rate is less dependent on the batch size
+        """
+        # start-snippet-2
+        # y.shape[0] is (symbolically) the number of rows in y, i.e.,
+        # number of examples (call it n) in the minibatch
+        # T.arange(y.shape[0]) is a symbolic vector which will contain
+        # [0,1,2,... n-1] T.log(self.p_y_given_x) is a matrix of
+        # Log-Probabilities (call it LP) with one row per example and
+        # one column per class LP[T.arange(y.shape[0]),y] is a vector
+        # v containing [LP[0,y[0]], LP[1,y[1]], LP[2,y[2]], ...,
+        # LP[n-1,y[n-1]]] and T.mean(LP[T.arange(y.shape[0]),y]) is
+        # the mean (across minibatch examples) of the elements in v,
+        # i.e., the mean log-likelihood across the minibatch.
+        return -T.mean(T.log(self.p_y_given_x)[T.arange(y.shape[0]), y])
+        # end-snippet-2
+
+    def errors(self, y):
+        """Return a float representing the number of errors in the minibatch
+        over the total number of examples of the minibatch ; zero one
+        loss over the size of the minibatch
+
+        :type y: theano.tensor.TensorType
+        :param y: corresponds to a vector that gives for each example the
+                  correct label
+        """
+
+        # check if y has same dimension of y_pred
+        if y.ndim != self.y_pred.ndim:
+            raise TypeError(
+                'y should have the same shape as self.y_pred',
+                ('y', y.type, 'y_pred', self.y_pred.type)
+            )
+        # check if y is of the correct datatype
+        if y.dtype.startswith('int'):
+            # the T.neq operator returns a vector of 0s and 1s, where 1
+            # represents a mistake in prediction
+            return T.mean(T.neq(self.y_pred, y))
+        else:
+            raise NotImplementedError()
+
+
+# start-snippet-1
+class HiddenLayer(object):
+    def __init__(self, rng, input, n_in, n_out, W=None, b=None,
+                 activation=T.tanh):
+        """
+        Typical hidden layer of a MLP: units are fully-connected and have
+        sigmoidal activation function. Weight matrix W is of shape (n_in,n_out)
+        and the bias vector b is of shape (n_out,).
+
+        NOTE : The nonlinearity used here is tanh
+
+        Hidden unit activation is given by: tanh(dot(input,W) + b)
+
+        :type rng: numpy.random.RandomState
+        :param rng: a random number generator used to initialize weights
+
+        :type input: theano.tensor.dmatrix
+        :param input: a symbolic tensor of shape (n_examples, n_in)
+
+        :type n_in: int
+        :param n_in: dimensionality of input
+
+        :type n_out: int
+        :param n_out: number of hidden units
+
+        :type activation: theano.Op or function
+        :param activation: Non linearity to be applied in the hidden
+                           layer
+        """
+        self.input = input
+        # end-snippet-1
+
+        # `W` is initialized with `W_values` which is uniformely sampled
+        # from sqrt(-6./(n_in+n_hidden)) and sqrt(6./(n_in+n_hidden))
+        # for tanh activation function
+        # the output of uniform if converted using asarray to dtype
+        # theano.config.floatX so that the code is runable on GPU
+        # Note : optimal initialization of weights is dependent on the
+        #        activation function used (among other things).
+        #        For example, results presented in [Xavier10] suggest that you
+        #        should use 4 times larger initial weights for sigmoid
+        #        compared to tanh
+        #        We have no info for other function, so we use the same as
+        #        tanh.
+        if W is None:
+            W_values = numpy.asarray(
+                rng.uniform(
+                    low=-numpy.sqrt(6. / (n_in + n_out)),
+                    high=numpy.sqrt(6. / (n_in + n_out)),
+                    size=(n_in, n_out)
+                ),
+                dtype=theano.config.floatX
+            )
+            if activation == theano.tensor.nnet.sigmoid:
+                W_values *= 4
+
+            W = theano.shared(value=W_values, name='W', borrow=True)
+
+        if b is None:
+            b_values = numpy.zeros((n_out,), dtype=theano.config.floatX)
+            b = theano.shared(value=b_values, name='b', borrow=True)
+
+        self.W = W
+        self.b = b
+
+        lin_output = T.dot(input, self.W) + self.b
+        self.output = (
+            lin_output if activation is None
+            else activation(lin_output)
+        )
+        # parameters of the model
+        self.params = [self.W, self.b]
+
+
+# start-snippet-2
+class MLP(object):
+    """Multi-Layer Perceptron Class
+
+    A multilayer perceptron is a feedforward artificial neural network model
+    that has one layer or more of hidden units and nonlinear activations.
+    Intermediate layers usually have as activation function tanh or the
+    sigmoid function (defined here by a ``HiddenLayer`` class)  while the
+    top layer is a softmax layer (defined here by a ``LogisticRegression``
+    class).
+    """
+
+    def __init__(self, rng, input, n_in, n_hidden, n_out):
+        """Initialize the parameters for the multilayer perceptron
+
+        :type rng: numpy.random.RandomState
+        :param rng: a random number generator used to initialize weights
+
+        :type input: theano.tensor.TensorType
+        :param input: symbolic variable that describes the input of the
+        architecture (one minibatch)
+
+        :type n_in: int
+        :param n_in: number of input units, the dimension of the space in
+        which the datapoints lie
+
+        :type n_hidden: int
+        :param n_hidden: number of hidden units
+
+        :type n_out: int
+        :param n_out: number of output units, the dimension of the space in
+        which the labels lie
+
+        """
+
+        # Since we are dealing with a one hidden layer MLP, this will translate
+        # into a HiddenLayer with a tanh activation function connected to the
+        # LogisticRegression layer; the activation function can be replaced by
+        # sigmoid or any other nonlinear function
+        self.hidden = HiddenLayer(
+            rng=rng,
+            input=input,
+            n_in=n_in,
+            n_out=n_hidden,
+            activation=T.tanh
+        )
+
+        # The logistic regression layer gets as input the hidden units
+        # of the hidden layer
+        self.maxent = LogisticRegression(
+            input=self.hidden.output,
+            n_in=n_hidden,
+            n_out=n_out
+        )
+        # L1 norm ; one regularization option is to enforce L1 norm to
+        # be small
+        self.L1 = abs(self.hidden.W).sum() + abs(self.maxent.W).sum()
+
+        # square of L2 norm ; one regularization option is to enforce
+        # square of L2 norm to be small
+        self.L2_sqr = (self.hidden.W ** 2).sum() + (self.maxent.W ** 2).sum()
+
+        # negative log likelihood of the MLP is given by the negative
+        # log likelihood of the output of the model, computed in the
+        # logistic regression layer
+        self.neg_ll = self.maxent.neg_ll
+        # same holds for the function computing the number of errors
+        self.errors = self.maxent.errors
+
+        # the parameters of the model are the parameters of the two layer it is
+        # made out of
+        self.params = self.hidden.params + self.maxent.params
+
+
+
+
+def test_mlp(learning_rate=0.01, L1_reg=0.00, L2_reg=0.0001, n_epochs=1000,
+             dataset='mnist.pkl.gz', batch_size=1, n_hidden=500):
+    """
+    Demonstrate stochastic gradient descent optimization for a multilayer
+    perceptron
+
+    This is demonstrated on MNIST.
+
+    :type learning_rate: float
+    :param learning_rate: learning rate used (factor for the stochastic
+    gradient
+
+    :type L1_reg: float
+    :param L1_reg: L1-norm's weight when added to the cost (see
+    regularization)
+
+    :type L2_reg: float
+    :param L2_reg: L2-norm's weight when added to the cost (see
+    regularization)
+
+    :type n_epochs: int
+    :param n_epochs: maximal number of epochs to run the optimizer
+
+    :type dataset: string
+    :param dataset: the path of the MNIST dataset file from
+                 http://www.iro.umontreal.ca/~lisa/deep/data/mnist/mnist.pkl.gz
+    """
+    datasets = load_data(dataset)
+
+    train_set_x, train_set_y = datasets[0]
+    valid_set_x, valid_set_y = datasets[1]
+    test_set_x, test_set_y = datasets[2]
+
+    ######################
+    # BUILD ACTUAL MODEL #
+    ######################
+    print '... building the model'
+
+    # allocate symbolic variables for the data
+    index = T.lscalar()  # index to a [mini]batch
+    x = T.matrix('x')  # the data is presented as rasterized images
+    y = T.ivector('y')  # the labels are presented as 1D vector of
+                        # [int] labels
+
+    rng = numpy.random.RandomState(1234)
+
+    # construct the MLP class
+    mlp = MLP(
+        rng=rng,
+        input=x,
+        n_in=28 * 28,
+        n_hidden=n_hidden,
+        n_out=10
+    )
+
+    # the cost we minimize during training is the negative log likelihood of
+    # the model plus the regularization terms (L1 and L2); cost is expressed
+    # here symbolically
+
+    # compiling a Theano function that computes the mistakes that are made
+    # by the model on a minibatch
+    test_model = theano.function(
+        inputs=[index],
+        outputs=mlp.maxent.errors(y),
+        givens={
+            x: test_set_x[index:index+1],
+            y: test_set_y[index:index+1]
+        }
+    )
+
+    validate_model = theano.function(
+        inputs=[index],
+        outputs=mlp.maxent.errors(y),
+        givens={
+            x: valid_set_x[index:index+1],
+            y: valid_set_y[index:index+1]
+        }
+    )
+
+    # compute the gradient of cost with respect to theta (sotred in params)
+    # the resulting gradients will be stored in a list gparams
+    cost = mlp.neg_ll(y) + L1_reg * mlp.L1 + L2_reg * mlp.L2_sqr
+    gparams = [T.grad(cost, param) for param in mlp.params]
+
+    # specify how to update the parameters of the model as a list of
+    # (variable, update expression) pairs
+
+    updates = [(mlp.params[i], mlp.params[i] - (learning_rate * gparams[i]))
+               for i in xrange(len(gparams))]
+
+    # compiling a Theano function `train_model` that returns the cost, but
+    # in the same time updates the parameter of the model based on the rules
+    # defined in `updates`
+    train_model = theano.function(
+        inputs=[index],
+        outputs=cost,
+        updates=updates,
+        givens={
+            x: train_set_x[index:index+1],
+            y: train_set_y[index:index+1]
+        }
+    )
+    # end-snippet-5
+
+    ###############
+    # TRAIN MODEL #
+    ###############
+    print '... training'
+
+    start_time = time.clock()
+
+    n_examples = train_set_x.get_value(borrow=True).shape[0]
+    n_dev_examples = valid_set_x.get_value(borrow=True).shape[0]
+    n_test_examples = test_set_x.get_value(borrow=True).shape[0]
+    
+    for epoch in range(1, n_epochs+1):
+        for idx in xrange(n_examples):
+            train_model(idx)
+        # compute zero-one loss on validation set
+        error = numpy.mean(map(validate_model, xrange(n_dev_examples)))
+        print('epoch %i, validation error %f %%' % (epoch, error * 100))
+
+    end_time = time.clock()
+    print >> sys.stderr, ('The code for file ' +
+                          os.path.split(__file__)[1] +
+                          ' ran for %.2fm' % ((end_time - start_time) / 60.))
+
+
+if __name__ == '__main__':
+    test_mlp()

spacy/_ml.pxd

@@ -5,6 +5,7 @@ from cymem.cymem cimport Pool
 from thinc.learner cimport LinearModel
 from thinc.features cimport Extractor, Feature
 from thinc.typedefs cimport atom_t, feat_t, weight_t, class_t
+from thinc.api cimport ExampleC
 
 from preshed.maps cimport PreshMapArray
 
@@ -14,9 +15,14 @@ from .tokens cimport Tokens
 
 cdef int arg_max(const weight_t* scores, const int n_classes) nogil
 
+cdef int arg_max_if_true(const weight_t* scores, const int* is_valid,  int n_classes) nogil
+
+cdef int arg_max_if_zero(const weight_t* scores, const int* costs, int n_classes) nogil
+
 
 cdef class Model:
-    cdef int n_classes
+    cdef readonly int n_classes
+    cdef readonly int n_feats
     
     cdef const weight_t* score(self, atom_t* context) except NULL
     cdef int set_scores(self, weight_t* scores, atom_t* context) except -1

spacy/_ml.pyx

@@ -2,6 +2,8 @@
 from __future__ import unicode_literals
 from __future__ import division
 
+from libc.string cimport memset
+
 from os import path
 import os
 import shutil
@@ -10,6 +12,7 @@ import cython
 import numpy.random
 
 from thinc.features cimport Feature, count_feats
+from thinc.api cimport Example
 
 
 cdef int arg_max(const weight_t* scores, const int n_classes) nogil:
@@ -23,17 +26,62 @@ cdef int arg_max(const weight_t* scores, const int n_classes) nogil:
     return best
 
 
+cdef int arg_max_if_true(const weight_t* scores, const int* is_valid,
+                         const int n_classes) nogil:
+    cdef int i
+    cdef int best = -1
+    cdef weight_t mode = 0
+    for i in range(n_classes):
+        if is_valid[i] and (best == -1 or scores[i] > mode):
+            mode = scores[i]
+            best = i
+    return best
+
+
+class ValidationError(Exception):
+    pass
+
+cdef int arg_max_if_zero(const weight_t* scores, const int* costs,
+                         const int n_classes) nogil:
+    cdef int i
+    cdef int best = -1
+    cdef weight_t mode = 0
+    for i in range(n_classes):
+        if costs[i] == 0 and (best == -1 or scores[i] > mode):
+            mode = scores[i]
+            best = i
+    return best
+
+
 cdef class Model:
     def __init__(self, n_classes, templates, model_loc=None):
         if model_loc is not None and path.isdir(model_loc):
             model_loc = path.join(model_loc, 'model')
         self.n_classes = n_classes
         self._extractor = Extractor(templates)
+        self.n_feats = self._extractor.n_templ
         self._model = LinearModel(n_classes, self._extractor.n_templ)
         self.model_loc = model_loc
         if self.model_loc and path.exists(self.model_loc):
             self._model.load(self.model_loc, freq_thresh=0)
 
+    def predict(self, Example eg):
+        assert self.n_classes == eg.c.nr_class
+        memset(eg.c.scores, 0, sizeof(weight_t) * eg.c.nr_class)
+        self.set_scores(eg.c.scores, eg.c.atoms)
+
+        eg.c.guess = arg_max_if_true(eg.c.scores, eg.c.is_valid, self.n_classes)
+        if eg.c.guess == -1:
+            raise ValidationError("No valid classes during prediction")
+
+    def train(self, Example eg):
+        self.predict(eg)
+        eg.c.best = arg_max_if_zero(eg.c.scores, eg.c.costs, self.n_classes)
+        if eg.c.best == -1:
+            raise ValidationError("No zero-cost classes during training.")
+        eg.c.cost = eg.c.costs[eg.c.guess]
+        self.update(eg.c.atoms, eg.c.guess, eg.c.best, eg.c.cost)
+
     cdef const weight_t* score(self, atom_t* context) except NULL:
         cdef int n_feats
         feats = self._extractor.get_feats(context, &n_feats)

spacy/_nn.py

@@ -0,0 +1,3 @@
+"""Feed-forward neural network, using Thenao."""
+
+

spacy/_nn.pyx

@@ -0,0 +1,146 @@
+"""Feed-forward neural network, using Thenao."""
+
+import os
+import sys
+import time
+
+import numpy
+
+import theano
+import theano.tensor as T
+import plac
+
+from spacy.gold import read_json_file
+from spacy.gold import GoldParse
+from spacy.en.pos import POS_TEMPLATES, POS_TAGS, setup_model_dir
+
+
+def build_model(n_classes, n_vocab, n_hidden, n_word_embed, n_tag_embed):
+    # allocate symbolic variables for the data
+    words = T.vector('words')
+    tags = T.vector('tags') 
+    
+    word_e = _init_embedding(n_words, n_word_embed)
+    tag_e = _init_embedding(n_tags, n_tag_embed)
+    label_e = _init_embedding(n_labels, n_label_embed)
+    maxent_W, maxent_b = _init_maxent_weights(n_hidden, n_classes)
+    hidden_W, hidden_b = _init_hidden_weights(28*28, n_hidden, T.tanh) 
+    params = [hidden_W, hidden_b, maxent_W, maxent_b, word_e, tag_e, label_e]
+
+    x = T.concatenate([
+          T.flatten(word_e[word_indices], outdim=1),
+          T.flatten(tag_e[tag_indices], outdim=1)])
+
+    p_y_given_x = feed_layer(
+                    T.nnet.softmax,
+                    maxent_W,
+                    maxent_b,
+                      feed_layer(
+                        T.tanh,
+                        hidden_W,
+                        hidden_b,
+                        x))[0]
+
+    guess = T.argmax(p_y_given_x)
+
+    cost = (
+        -T.log(p_y_given_x[y])
+        + L1(L1_reg, maxent_W, hidden_W, word_e, tag_e)
+        + L2(L2_reg, maxent_W, hidden_W, wod_e, tag_e)
+    )
+
+    train_model = theano.function(
+        inputs=[words, tags, y],
+        outputs=guess,
+        updates=[update(learning_rate, param, cost) for param in params]
+    )
+
+    evaluate_model = theano.function(
+        inputs=[x, y],
+        outputs=T.neq(y, T.argmax(p_y_given_x[0])),
+    )
+    return train_model, evaluate_model
+
+
+def _init_embedding(vocab_size, n_dim):
+    embedding = 0.2 * numpy.random.uniform(-1.0, 1.0, (vocab_size+1, n_dim))
+    return theano.shared(embedding).astype(theano.config.floatX)
+
+
+def _init_maxent_weights(n_hidden, n_out):
+    weights = numpy.zeros((n_hidden, 10), dtype=theano.config.floatX)
+    bias =  numpy.zeros((10,), dtype=theano.config.floatX)
+    return (
+        theano.shared(name='W', borrow=True, value=weights),
+        theano.shared(name='b', borrow=True, value=bias)
+    )
+
+
+def _init_hidden_weights(n_in, n_out, activation=T.tanh):
+    rng = numpy.random.RandomState(1234)
+    weights = numpy.asarray(
+        rng.uniform(
+            low=-numpy.sqrt(6. / (n_in + n_out)),
+            high=numpy.sqrt(6. / (n_in + n_out)),
+            size=(n_in, n_out)
+        ),
+        dtype=theano.config.floatX
+    )
+
+    bias = numpy.zeros((n_out,), dtype=theano.config.floatX)
+    return (
+        theano.shared(value=weights, name='W', borrow=True),
+        theano.shared(value=bias, name='b', borrow=True)
+    )
+
+           
+def feed_layer(activation, weights, bias, input):
+    return activation(T.dot(input, weights) + bias)
+
+
+def L1(L1_reg, w1, w2):
+    return L1_reg * (abs(w1).sum() + abs(w2).sum())
+
+
+def L2(L2_reg, w1, w2):
+    return L2_reg * ((w1 ** 2).sum() + (w2 ** 2).sum())
+
+
+def update(eta, param, cost):
+    return (param, param - (eta * T.grad(cost, param)))
+
+
+def main(train_loc, eval_loc, model_dir):
+    learning_rate = 0.01
+    L1_reg = 0.00
+    L2_reg = 0.0001
+
+    print "... reading the data"
+    gold_train = list(read_json_file(train_loc))
+    print '... building the model'
+    pos_model_dir = path.join(model_dir, 'pos')
+    if path.exists(pos_model_dir):
+        shutil.rmtree(pos_model_dir)
+    os.mkdir(pos_model_dir)
+
+    setup_model_dir(sorted(POS_TAGS.keys()), POS_TAGS, POS_TEMPLATES, pos_model_dir)
+
+    train_model, evaluate_model = build_model(n_hidden, len(POS_TAGS), learning_rate,
+                                              L1_reg, L2_reg)
+
+    print '... training'
+    for epoch in range(1, n_epochs+1):
+        for raw_text, sents in gold_tuples:
+            for (ids, words, tags, ner, heads, deps), _ in sents:
+                tokens = nlp.tokenizer.tokens_from_list(words)
+                for t in tokens:
+                    guess = train_model([t.orth], [t.tag])
+                    loss += guess != t.tag
+        print loss
+        # compute zero-one loss on validation set
+        #error = numpy.mean([evaluate_model(x, y) for x, y in dev_examples])
+        #print('epoch %i, validation error %f %%' % (epoch, error * 100))
+
+
+if __name__ == '__main__':
+    plac.call(main)

spacy/_theano.pxd

@@ -0,0 +1,13 @@
+from ._ml cimport Model
+from thinc.nn cimport InputLayer
+
+
+cdef class TheanoModel(Model):
+    cdef InputLayer input_layer
+    cdef object train_func
+    cdef object predict_func
+    cdef object debug
+
+    cdef public float eta
+    cdef public float mu
+    cdef public float t

spacy/_theano.pyx

@@ -0,0 +1,52 @@
+from thinc.api cimport Example, ExampleC
+from thinc.typedefs cimport weight_t
+
+from ._ml cimport arg_max_if_true
+from ._ml cimport arg_max_if_zero
+
+import numpy
+from os import path
+
+
+cdef class TheanoModel(Model):
+    def __init__(self, n_classes, input_spec, train_func, predict_func, model_loc=None,
+                 eta=0.001, mu=0.9, debug=None):
+        if model_loc is not None and path.isdir(model_loc):
+            model_loc = path.join(model_loc, 'model')
+
+        self.eta = eta
+        self.mu = mu
+        self.t = 1
+        initializer = lambda: 0.2 * numpy.random.uniform(-1.0, 1.0)
+        self.input_layer = InputLayer(input_spec, initializer)
+        self.train_func = train_func
+        self.predict_func = predict_func
+        self.debug = debug
+
+        self.n_classes = n_classes
+        self.n_feats = len(self.input_layer)
+        self.model_loc = model_loc
+        
+    def predict(self, Example eg):
+        self.input_layer.fill(eg.embeddings, eg.atoms, use_avg=True)
+        theano_scores = self.predict_func(eg.embeddings)[0]
+        cdef int i
+        for i in range(self.n_classes):
+            eg.c.scores[i] = theano_scores[i]
+        eg.c.guess = arg_max_if_true(eg.c.scores, eg.c.is_valid, self.n_classes)
+
+    def train(self, Example eg):
+        self.input_layer.fill(eg.embeddings, eg.atoms, use_avg=False)
+        theano_scores, update, y, loss = self.train_func(eg.embeddings, eg.costs,
+                                                         self.eta, self.mu)
+        self.input_layer.update(update, eg.atoms, self.t, self.eta, self.mu)
+        for i in range(self.n_classes):
+            eg.c.scores[i] = theano_scores[i]
+        eg.c.guess = arg_max_if_true(eg.c.scores, eg.c.is_valid, self.n_classes)
+        eg.c.best = arg_max_if_zero(eg.c.scores, eg.c.costs, self.n_classes)
+        eg.c.cost = eg.c.costs[eg.c.guess]
+        eg.c.loss = loss
+        self.t += 1
+
+    def end_training(self):
+        pass

spacy/en/__init__.py

@@ -13,6 +13,9 @@ from ..multi_words import RegexMerger
 
 from .pos import EnPosTagger
 from .pos import POS_TAGS
+
+from ..wsd.supersense_tagger import SenseTagger
+
 from .attrs import get_flags
 from . import regexes
 
@@ -80,6 +83,7 @@ class English(object):
             self.has_parser_model = False
             self.has_tagger_model = False
             self.has_entity_model = False
+            self.has_senser_model = False
         else:
             tok_data_dir = path.join(data_dir, 'tokenizer')
             tok_rules, prefix_re, suffix_re, infix_re = read_lang_data(tok_data_dir)
@@ -89,6 +93,7 @@ class English(object):
             self.has_parser_model = path.exists(path.join(self._data_dir, 'deps'))
             self.has_tagger_model = path.exists(path.join(self._data_dir, 'pos'))
             self.has_entity_model = path.exists(path.join(self._data_dir, 'ner'))
+            self.has_senser_model = path.exists(path.join(self._data_dir, 'wsd'))
 
         self.tokenizer = Tokenizer(self.vocab, tok_rules, prefix_re,
                                    suffix_re, infix_re,
@@ -102,6 +107,7 @@ class English(object):
         self._tagger = None
         self._parser = None
         self._entity = None
+        self._senser = None
 
     @property
     def tagger(self):
@@ -109,6 +115,12 @@ class English(object):
             self._tagger = EnPosTagger(self.vocab.strings, self._data_dir)
         return self._tagger
 
+    @property
+    def senser(self):
+        if self._senser is None:
+            self._senser = SenseTagger(self.vocab.strings, self._data_dir)
+        return self._senser
+
     @property
     def parser(self):
         if self._parser is None:

spacy/gold.pxd

@@ -1,6 +1,8 @@
 from cymem.cymem cimport Pool
 
 from .structs cimport TokenC
+from .typedefs cimport flags_t
+
 from .syntax.transition_system cimport Transition
 
 cimport numpy
@@ -10,6 +12,7 @@ cdef struct GoldParseC:
     int* tags
     int* heads
     int* labels
+    flags_t* ssenses
     int** brackets
     Transition* ner
 
@@ -25,6 +28,7 @@ cdef class GoldParse:
     cdef readonly list heads
     cdef readonly list labels
     cdef readonly dict orths
+    cdef readonly list ssenses
     cdef readonly list ner
     cdef readonly list ents
     cdef readonly dict brackets

spacy/gold.pyx

@@ -9,6 +9,8 @@ from os import path
 
 from libc.string cimport memset
 
+from .typedefs cimport flags_t
+
 
 def tags_to_entities(tags):
     entities = []
@@ -202,6 +204,7 @@ cdef class GoldParse:
         self.c.tags = <int*>self.mem.alloc(len(tokens), sizeof(int))
         self.c.heads = <int*>self.mem.alloc(len(tokens), sizeof(int))
         self.c.labels = <int*>self.mem.alloc(len(tokens), sizeof(int))
+        self.c.ssenses = <flags_t*>self.mem.alloc(len(tokens), sizeof(flags_t))
         self.c.ner = <Transition*>self.mem.alloc(len(tokens), sizeof(Transition))
         self.c.brackets = <int**>self.mem.alloc(len(tokens), sizeof(int*))
         for i in range(len(tokens)):
@@ -211,15 +214,20 @@ cdef class GoldParse:
         self.heads = [None] * len(tokens)
         self.labels = [''] * len(tokens)
         self.ner = ['-'] * len(tokens)
+        self.ssenses = [[] for _ in range(len(tokens))]
 
         self.cand_to_gold = align([t.orth_ for t in tokens], annot_tuples[1])
         self.gold_to_cand = align(annot_tuples[1], [t.orth_ for t in tokens])
 
         self.orig_annot = zip(*annot_tuples)
 
+        # This iterates 0...n for n words in the candidate, with an index
+        # gold_i aligned into the gold. Assign tag, label, ner and word sense.
+        # For the head, the value is an index into the gold sentence, so we
+        # have to translate it across into the candidate.
         for i, gold_i in enumerate(self.cand_to_gold):
             if gold_i is None:
-                # TODO: What do we do for missing values again?
+                # Missing values handled in the various oracle functions
                 pass
             else:
                 self.tags[i] = annot_tuples[2][gold_i]
@@ -244,6 +252,7 @@ cdef class GoldParse:
                                 self.labels[w1] = ''
                                 self.heads[w2] = None
                                 self.labels[w2] = ''
+                                self.ssenses[w2] = []
 
         # Check there are no cycles in the dependencies, i.e. we are a tree
         for w in range(self.length):

spacy/lexeme.pyx

@@ -28,9 +28,12 @@ cdef int set_lex_struct_props(LexemeC* lex, dict props, StringStore string_store
     lex.sentiment = props['sentiment']
 
     lex.flags = props['flags']
-    cdef flags_t sense_id
-    for sense_id in props.get('senses', []):
-        lex.senses |= 1 << sense_id
+    cdef flags_t sense_id = 0
+    cdef flags_t one = 1
+    lex.senses = 0
+    for _sense_id in props.get('senses', []):
+        sense_id = _sense_id
+        lex.senses |= one << sense_id
     lex.repvec = empty_vec
 
 
@@ -48,7 +51,9 @@ cdef class Lexeme:
         return self.l2_norm != 0
 
     cpdef bint check(self, attr_id_t flag_id) except -1:
-        return self.flags & (1 << flag_id)
+        cdef flags_t one = 1
+        return self.flags & (one << flag_id)
 
     cpdef bint has_sense(self, flags_t flag_id) except -1:
-        return self.senses & (1 << flag_id)
+        cdef flags_t one = 1
+        return self.senses & (one << flag_id)

spacy/munge/corpus.py

@@ -0,0 +1,100 @@
+from __future__ import unicode_literals
+import zlib
+import gzip
+import sqlite3
+import os
+from os import path
+import sys
+
+
+class DocsDB(object):
+    def __init__(self, db_loc, batch_size=1000, limit=1000):
+        limit = int(limit)
+        batch_size = int(batch_size)
+        self._conn = sqlite3.connect(db_loc)
+        self._curr = self._conn.cursor()
+        try:
+            self._curr.execute('SELECT * FROM docs')
+        except:
+            print db_loc
+            raise
+        self._batch = self._curr.fetchmany(batch_size)
+        self._batch_size = batch_size
+        self._limit = limit
+
+    def __iter__(self):
+        while self._batch:
+            for doc_id, compressed_doc in self._batch:
+                if compressed_doc:
+                    yield zlib.decompress(compressed_doc).decode('ascii')
+                if doc_id >= self._limit:
+                    self._batch = None
+                    break
+            else:
+                self._batch = self._curr.fetchmany(size=self._batch_size)
+
+
+class Gigaword(DocsDB):
+    @classmethod
+    def create(cls, giga_dir, db_loc):
+        giga_dir = str(giga_dir)
+        db_loc = str(db_loc)
+        if path.exists(db_loc):
+            os.unlink(db_loc)
+        conn = sqlite3.connect(db_loc)
+        c = conn.cursor()
+        c.execute('''CREATE TABLE docs (id INTEGER PRIMARY KEY, body BLOB)''')
+        doc_id = 0
+        for file_loc in iter_files(giga_dir):
+            print >> sys.stderr, file_loc
+            for doc in iter_docs(file_loc):
+                if doc.strip():
+                    compressed = sqlite3.Binary(zlib.compress(doc))
+                    c.execute('''INSERT INTO docs VALUES (?, ?)''', (doc_id, compressed))
+                    doc_id += 1
+            conn.commit()
+
+
+def iter_files(giga_dir):
+    for subdir in os.listdir(giga_dir):
+        if not path.isdir(path.join(giga_dir, subdir)):
+            continue
+        for filename in os.listdir(path.join(giga_dir, subdir)):
+            if filename.endswith('gz'):
+                yield path.join(giga_dir, subdir, filename)
+ 
+
+def iter_docs(zip_loc):
+    doc = []
+    para = []
+    in_doc = False
+    in_para = False
+    try:
+        lines = gzip.open(zip_loc, 'r').read().replace('&', '&').split('\n')
+    except UnicodeDecodeError:
+        lines = []
+    for line in lines:
+        line = line.strip()
+        if not line:
+            pass
+        elif line[0] != '<':
+            if in_para:
+                para.append(line)
+        elif line.startswith('<DOC'):
+            in_doc = True
+        elif line.startswith('<P>'):
+            assert in_doc
+            in_para = True
+        elif line.startswith('</DOC'):
+            assert not in_para
+            in_doc = False
+            yield '\n\n'.join(doc)
+            doc = []
+        elif line.startswith('</P>'):
+            doc.append(' '.join(para))
+            in_para = False
+            para = []
+        else:
+            pass
+    assert not in_doc
+    assert not in_para

spacy/munge/read_semcor.py

@@ -0,0 +1,193 @@
+from __future__ import unicode_literals
+from __future__ import division
+import plac
+import re
+from os import path
+import os
+import codecs
+
+from spacy.en import English
+
+lexnames_str = """
+-1      NO_SENSE       -1
+00      J_all 3
+01      A_pert        3 
+02      A_all 4
+03      N_Tops       1  
+04      N_act        1
+05      N_animal     1
+06      N_artifact   1
+07      N_attribute  1
+08      N_body       1
+09      N_cognition  1
+10      N_communication      1
+11      N_event      1
+12      N_feeling    1
+13      N_food       1
+14      N_group      1
+15      N_location   1
+16      N_motive     1
+17      N_object     1
+18      N_person     1
+19      N_phenomenon 1
+20      N_plant      1
+21      N_possession 1
+22      N_process    1
+23      N_quantity   1
+24      N_relation   1
+25      N_shape      1
+26      N_state      1
+27      N_substance  1
+28      N_time       1
+29      V_body       2
+30      V_change     2
+31      V_cognition  2
+32      V_communication      2
+33      V_competition        2
+34      V_consumption        2
+35      V_contact    2
+36      V_creation   2
+37      V_emotion    2
+38      V_motion     2
+39      V_perception 2
+40      V_possession 2
+41      V_social     2
+42      V_stative    2
+43      V_weather    2
+44      A_ppl 3
+""".strip()
+
+SUPERSENSES = tuple(line.split()[1] for line in lexnames_str.split('\n'))
+
+
+
+
+def re_get(exp, string):
+    obj = exp.search(string)
+    if obj is None:
+        return obj
+    else:
+        return obj.group()
+
+
+lemma_re = re.compile(r'(?<=lemma=)[^ >]+')
+cmd_re = re.compile(r'(?<=cmd=)[^ >]+')
+pos_re = re.compile(r'(?<=pos=)[^ >]+')
+ot_re = re.compile(r'(?<=ot=)[^ >]+')
+wnsn_re = re.compile(r'(?<=wnsn=)[^ >]+')
+lexsn_re = re.compile(r'(?<=lexsn=)[^ >]+')
+supersense_re = re.compile(r'(?<=lexsn=\d:)\d\d')
+orth_re = re.compile(r'(?<=>)[^<]+(?=<)')
+class Token(object):
+    def __init__(self, line):
+        self.cmd = re_get(cmd_re, line)
+        self.lemma = re_get(lemma_re, line)
+        self.ot = re_get(ot_re, line)
+        self.pos = re_get(pos_re, line)
+        self.wnsn = re_get(wnsn_re, line)
+        self.lexsn = re_get(lexsn_re, line)
+        supersense = re_get(supersense_re, line)
+        if supersense is None:
+            self.supersense = SUPERSENSES[0]
+        else:
+            self.supersense = SUPERSENSES[int(supersense) + 1]
+        self.orth = re_get(orth_re, line)
+
+    def __str__(self):
+        return (self.cmd, self.lemma, self.ot, self.pos,
+                self.wnsn, self.lexsn, self.orth)
+
+    def __repr__(self):
+        return str(self)
+
+
+def read_file(loc):
+    paras = []
+    sents = []
+    sent = []
+    filename = None
+    pnum = None
+    snum = None
+    for line in codecs.open(loc, 'r', 'latin1'):
+        line = line.strip()
+        if not line:
+            continue
+
+        if line.startswith('contextfile'):
+            continue
+
+        if line.startswith('<context '):
+            assert filename is None
+            pieces = line.split()
+            filename = pieces[1].replace('filename=', '')
+            continue
+        
+        if line.startswith('<p '):
+            assert pnum is None
+            pnum = int(line.split('=')[1][:-1])
+            continue
+        
+        if line.startswith('<s '):
+            assert snum is None, line
+            snum = int(line.split('=')[1][:-1])
+            continue
+
+        if line.startswith('<wf ') or line.startswith('<punc'):
+            sent.append(Token(line))
+            continue
+
+        if line == '</s>':
+            sents.append((snum, sent))
+            sent = []
+            snum = None
+            continue
+
+        if line == '</p>':
+            paras.append((pnum, sents))
+            sents = []
+            pnum = None
+            continue
+    return paras
+
+
+def read_semcor(semcor_dir):
+    docs = []
+    brown1 = path.join(semcor_dir, 'brown1', 'tagfiles')
+    for filename in os.listdir(brown1):
+        file_path = path.join(brown1, filename)
+        docs.append((filename, read_file(file_path)))
+    return docs
+
+
+def test_token():
+    string = '<wf cmd=done pos=NN lemma=sheriff wnsn=1 lexsn=1:18:00::>sheriff</wf>'
+    token = Token(string)
+    assert token.cmd == 'done'
+    assert token.pos == 'NN'
+    assert token.lemma == 'sheriff'
+    assert token.wnsn == '1'
+    assert token.lexsn == '1:18:00::'
+    assert token.orth == 'sheriff'
+
+
+def main(model_dir, semcor_dir):
+    brown1 = path.join(semcor_dir, 'brown1', 'tagfiles')
+
+    nlp = English(data_dir=model_dir)
+    total_right = 0
+    total_wrong = 0
+    total_multi = 0
+    for filename in os.listdir(brown1):
+        file_path = path.join(brown1, filename)
+        annotations = read_file(file_path)
+
+        n_multi, n_right, n_wrong = eval_text(nlp, annotations)
+        total_right += n_right
+        total_wrong += n_wrong
+        total_multi += n_multi
+    print total_right, total_wrong
+    print total_right / (total_right + total_wrong)
+    print total_multi / (total_multi + total_right + total_wrong) 
+
+if __name__ == '__main__':
+    plac.call(main)

spacy/munge/read_wordnet.py

@@ -0,0 +1,115 @@
+"""
+Get a mapping of (lemma, sense_number)-->supersense, and a mapping
+(lemma, ON group)-->(lemma, sense_number).
+
+Then we can read the OntoNotes token-->(lemma, ON group) annotations, and resolve
+them to the token-->supersense annotations we want to train from.
+
+supersense: A WordNet lexical file number
+sense_number: A WordNet sense key, found in e.g. wordnet/index.sense file
+lex_filenum: A WordNet "super sense", or lexical file number.
+onto_group: An OntoNotes sense grouping, which dominates zero or more WN senses.
+"""
+from __future__ import division
+
+from os import path
+import os
+import re
+import codecs
+
+
+def get_sense_to_ssense(index_dot_sense_loc):
+    mapping = {}
+    pos_tags = [None, 'n', 'v', 'j', 'a', 's']
+    for line in codecs.open(index_dot_sense_loc, 'r', 'utf8'):
+        sense_key, synset_offset, sense_number, tag_cnt = line.split()
+        lemma, lex_sense = sense_key.split('%')
+        ss_type, lex_filenum, lex_id, head_word, head_id = lex_sense.split(':')
+        pos = pos_tags[int(ss_type)]
+        mapping[(lemma, pos, int(sense_number))] = int(lex_filenum)
+    return mapping
+
+
+sense_group_re = re.compile(r'<sense .*?</sense>', re.DOTALL)
+wn_mapping_re = re.compile(r'version="3.0">([^<]+)<')
+def get_og_to_sense(sense_inv_dir):
+    mapping = {}
+    for filename in os.listdir(sense_inv_dir):
+        if not filename.endswith('.xml'):
+            continue
+        if '-' not in filename:
+            continue
+        lemma, pos = filename.split('-')[:2]
+        pos = pos[0]
+        # Word is these often don't validate, because of course. So, just parse
+        # with regex...
+        xml_str = open(path.join(sense_inv_dir, filename)).read()
+        for sense_grouping in sense_group_re.findall(xml_str):
+            group_num = sense_grouping.split('n="')[1].split('"')[0]
+            if not group_num:
+                continue
+
+            group_num = int(float(group_num))
+            key = (lemma, pos, int(group_num))
+            mapping.setdefault(key, [])
+            wn_elem = wn_mapping_re.search(sense_grouping)
+            if wn_elem is not None:
+                sense_num_str = wn_elem.groups()[0].replace('.', ',')
+                sense_ids = [(lemma, pos, int(n)) for n in sense_num_str.strip().split(',')]
+                mapping[key].extend(sense_ids)
+    return mapping
+
+
+def get_lexnames(loc):
+    names = {}
+    for line in open(loc):
+        id_, name, syn_type = line.split()
+        names[int(id_)] = name
+    return names
+
+
+def get_og_to_ssenses(wordnet_dir, onto_dir):
+    sense_inv_dir = path.join(onto_dir, 'data', 'english', 'metadata', 'sense-inventories')
+    og_to_sense = get_og_to_sense(sense_inv_dir)
+    sense_to_ssense = get_sense_to_ssense(path.join(wordnet_dir, 'index.sense'))
+    lexnames = get_lexnames(path.join(wordnet_dir, 'lexnames'))
+
+    mapping = {}
+    for key, senses in og_to_sense.items():
+        if senses is not None:
+            mapping[key] = set([lexnames[sense_to_ssense[s_key]]
+                                for s_key in senses if s_key in sense_to_ssense])
+    return mapping
+
+
+def make_supersense_dict(wordnet_dir):
+    sense_to_ssense = get_sense_to_ssense(path.join(wordnet_dir, 'index.sense'))
+    gather = {}
+    for (word, pos, sense), supersense in sense_to_ssense.items():
+        key = (word, pos)
+        gather.setdefault((word, pos), []).append((int(sense), supersense))
+    mapping = {}
+    for (word, pos), senses in gather.items():
+        n_senses = len(senses)
+        probs = {}
+        remaining = 1.0
+        for sense, supersense in sorted(senses):
+            remaining /= 2
+            probs[supersense] = probs.get(supersense, 0.0) + remaining
+        for sense, supersense in sorted(senses):
+            probs[supersense] += remaining / len(senses)
+        mapping.setdefault(word, {}).update(probs)
+    return mapping
+
+
+def main(wordnet_dir, onto_dir):
+    mapping = make_supersense_dict(wordnet_dir)
+    print mapping[('dog', 'v')]
+    print mapping[('dog', 'n')]
+    print mapping[('abandon', 'v')]
+    print mapping[('abandon', 'n')]
+
+
+if __name__ == '__main__':
+    import plac
+    plac.call(main)

spacy/scorer.py

@@ -2,6 +2,8 @@ from __future__ import division
 
 from .gold import tags_to_entities
 
+from .senses import STRINGS as SENSE_STRINGS
+
 
 class PRFScore(object):
     """A precision / recall / F score"""
@@ -38,6 +40,7 @@ class Scorer(object):
         self.labelled = PRFScore()
         self.tags = PRFScore()
         self.ner = PRFScore()
+        self.wsd = PRFScore()
         self.eval_punct = eval_punct
 
     @property
@@ -73,7 +76,7 @@ class Scorer(object):
 
         gold_deps = set()
         gold_tags = set()
-        gold_ents = set(tags_to_entities([annot[-1] for annot in gold.orig_annot]))
+        gold_ents = set(tags_to_entities([annot[5] for annot in gold.orig_annot]))
         for id_, word, tag, head, dep, ner in gold.orig_annot:
             gold_tags.add((id_, tag))
             if dep.lower() not in ('p', 'punct'):

spacy/senses.pyx

@@ -1,88 +0,0 @@
-from __future__ import unicode_literals
-cimport parts_of_speech
-
-
-POS_SENSES[<int>parts_of_speech.NO_TAG] = 0
-POS_SENSES[<int>parts_of_speech.ADJ] = 0
-POS_SENSES[<int>parts_of_speech.ADV] = 0
-POS_SENSES[<int>parts_of_speech.ADP] = 0
-POS_SENSES[<int>parts_of_speech.CONJ] = 0
-POS_SENSES[<int>parts_of_speech.DET] = 0
-POS_SENSES[<int>parts_of_speech.NOUN] = 0
-POS_SENSES[<int>parts_of_speech.NUM] = 0
-POS_SENSES[<int>parts_of_speech.PRON] = 0
-POS_SENSES[<int>parts_of_speech.PRT] = 0
-POS_SENSES[<int>parts_of_speech.VERB] = 0
-POS_SENSES[<int>parts_of_speech.X] = 0
-POS_SENSES[<int>parts_of_speech.PUNCT] = 0
-POS_SENSES[<int>parts_of_speech.EOL] = 0
-
-
-cdef int _sense = 0
-
-for _sense in range(A_behavior, N_act):
-    POS_SENSES[<int>parts_of_speech.ADJ] |= 1 << _sense
-
-for _sense in range(N_act, V_body):
-    POS_SENSES[<int>parts_of_speech.NOUN] |= 1 << _sense
-
-for _sense in range(V_body, V_weather+1):
-    POS_SENSES[<int>parts_of_speech.VERB] |= 1 << _sense
-
-
-
-STRINGS = (
-    'A_behavior',
-    'A_body',
-    'A_feeling',
-    'A_mind',
-    'A_motion',
-    'A_perception',
-    'A_quantity',
-    'A_relation',
-    'A_social',
-    'A_spatial',
-    'A_substance',
-    'A_time',
-    'A_weather',
-    'N_act',
-    'N_animal',
-    'N_artifact',
-    'N_attribute',
-    'N_body',
-    'N_cognition',
-    'N_communication',
-    'N_event',
-    'N_feeling',
-    'N_food',
-    'N_group',
-    'N_location',
-    'N_motive',
-    'N_object',
-    'N_person',
-    'N_phenomenon',
-    'N_plant',
-    'N_possession',
-    'N_process',
-    'N_quantity',
-    'N_relation',
-    'N_shape',
-    'N_state',
-    'N_substance',
-    'N_time',
-    'V_body',
-    'V_change',
-    'V_cognition',
-    'V_communication',
-    'V_competition',
-    'V_consumption',
-    'V_contact',
-    'V_creation',
-    'V_emotion',
-    'V_motion',
-    'V_perception',
-    'V_possession',
-    'V_social',
-    'V_stative',
-    'V_weather'
-)

spacy/syntax/_parse_features.pyx

@@ -61,7 +61,7 @@ cdef inline void fill_token(atom_t* context, const TokenC* token) nogil:
         context[9] = token.lex.shape
         context[10] = token.ent_iob
         context[11] = token.ent_type
-        context[12] = token.lex.senses & senses.POS_SENSES[<int>token.pos]
+        context[12] = 0 # token.lex.senses & senses.POS_SENSES[<int>token.pos]
 
 cdef int fill_context(atom_t* ctxt, StateClass st) nogil:
     # Take care to fill every element of context!

spacy/syntax/arc_eager.pyx

@@ -398,7 +398,8 @@ cdef class ArcEager(TransitionSystem):
             n_valid += output[i]
         assert n_valid >= 1
 
-    cdef int set_costs(self, int* output, StateClass stcls, GoldParse gold) except -1:
+    cdef int set_costs(self, bint* is_valid, int* costs, 
+                       StateClass stcls, GoldParse gold) except -1:
         cdef int i, move, label
         cdef label_cost_func_t[N_MOVES] label_cost_funcs
         cdef move_cost_func_t[N_MOVES] move_cost_funcs
@@ -423,30 +424,14 @@ cdef class ArcEager(TransitionSystem):
         n_gold = 0
         for i in range(self.n_moves):
             if self.c[i].is_valid(stcls, self.c[i].label):
+                is_valid[i] = True
                 move = self.c[i].move
                 label = self.c[i].label
                 if move_costs[move] == -1:
                     move_costs[move] = move_cost_funcs[move](stcls, &gold.c)
-                output[i] = move_costs[move] + label_cost_funcs[move](stcls, &gold.c, label)
-                n_gold += output[i] == 0
+                costs[i] = move_costs[move] + label_cost_funcs[move](stcls, &gold.c, label)
+                n_gold += costs[i] == 0
             else:
-                output[i] = 9000
+                is_valid[i] = False
+                costs[i] = 9000
         assert n_gold >= 1
-
-    cdef Transition best_valid(self, const weight_t* scores, StateClass stcls) except *:
-        cdef bint[N_MOVES] is_valid
-        is_valid[SHIFT] = Shift.is_valid(stcls, -1)
-        is_valid[REDUCE] = Reduce.is_valid(stcls, -1)
-        is_valid[LEFT] = LeftArc.is_valid(stcls, -1)
-        is_valid[RIGHT] = RightArc.is_valid(stcls, -1)
-        is_valid[BREAK] = Break.is_valid(stcls, -1)
-        cdef Transition best
-        cdef weight_t score = MIN_SCORE
-        cdef int i
-        for i in range(self.n_moves):
-            if scores[i] > score and is_valid[self.c[i].move]:
-                best = self.c[i]
-                score = scores[i]
-        assert best.clas < self.n_moves
-        assert score > MIN_SCORE, (stcls.stack_depth(), stcls.buffer_length(), stcls.is_final(), stcls._b_i, stcls.length)
-        return best

spacy/syntax/joint.pxd

@@ -0,0 +1,17 @@
+from cymem.cymem cimport Pool
+
+from thinc.typedefs cimport weight_t
+
+from .stateclass cimport StateClass
+
+from .transition_system cimport TransitionSystem, Transition
+from ..gold cimport GoldParseC
+
+
+cdef class ArcEager(TransitionSystem):
+    pass
+
+
+cdef int push_cost(StateClass stcls, const GoldParseC* gold, int target) nogil
+cdef int arc_cost(StateClass stcls, const GoldParseC* gold, int head, int child) nogil
+

spacy/syntax/joint.pyx

@@ -0,0 +1,452 @@
+# cython: profile=True
+from __future__ import unicode_literals
+
+import ctypes
+import os
+
+from ..structs cimport TokenC
+
+from .transition_system cimport do_func_t, get_cost_func_t
+from .transition_system cimport move_cost_func_t, label_cost_func_t
+from ..gold cimport GoldParse
+from ..gold cimport GoldParseC
+
+from libc.stdint cimport uint32_t
+from libc.string cimport memcpy
+
+from cymem.cymem cimport Pool
+from .stateclass cimport StateClass
+
+
+DEF NON_MONOTONIC = True
+DEF USE_BREAK = True
+DEF USE_ROOT_ARC_SEGMENT = True
+
+cdef weight_t MIN_SCORE = -90000
+
+# Break transition from here
+# http://www.aclweb.org/anthology/P13-1074
+cdef enum:
+    SHIFT
+    REDUCE
+    LEFT
+    RIGHT
+
+    BREAK
+
+    N_MOVES
+
+
+MOVE_NAMES = [None] * N_MOVES
+MOVE_NAMES[SHIFT] = 'S'
+MOVE_NAMES[REDUCE] = 'D'
+MOVE_NAMES[LEFT] = 'L'
+MOVE_NAMES[RIGHT] = 'R'
+MOVE_NAMES[BREAK] = 'B'
+
+
+# Helper functions for the arc-eager oracle
+
+cdef int push_cost(StateClass stcls, const GoldParseC* gold, int target) nogil:
+    cdef int cost = 0
+    cdef int i, S_i
+    for i in range(stcls.stack_depth()):
+        S_i = stcls.S(i)
+        if gold.heads[target] == S_i:
+            cost += 1
+        if gold.heads[S_i] == target and (NON_MONOTONIC or not stcls.has_head(S_i)):
+            cost += 1
+    cost += Break.is_valid(stcls, -1) and Break.move_cost(stcls, gold) == 0
+    return cost
+
+
+cdef int pop_cost(StateClass stcls, const GoldParseC* gold, int target) nogil:
+    cdef int cost = 0
+    cdef int i, B_i
+    for i in range(stcls.buffer_length()):
+        B_i = stcls.B(i)
+        cost += gold.heads[B_i] == target
+        cost += gold.heads[target] == B_i
+        if gold.heads[B_i] == B_i or gold.heads[B_i] < target:
+            break
+    cost += Break.is_valid(stcls, -1) and Break.move_cost(stcls, gold) == 0
+    return cost
+
+
+cdef int arc_cost(StateClass stcls, const GoldParseC* gold, int head, int child) nogil:
+    if arc_is_gold(gold, head, child):
+        return 0
+    elif stcls.H(child) == gold.heads[child]:
+        return 1
+    # Head in buffer
+    elif gold.heads[child] >= stcls.B(0) and stcls.B(1) != -1:
+        return 1
+    else:
+        return 0
+
+
+cdef bint arc_is_gold(const GoldParseC* gold, int head, int child) nogil:
+    if gold.labels[child] == -1:
+        return True
+    elif USE_ROOT_ARC_SEGMENT and _is_gold_root(gold, head) and _is_gold_root(gold, child):
+        return True
+    elif gold.heads[child] == head:
+        return True
+    else:
+        return False
+
+
+cdef bint label_is_gold(const GoldParseC* gold, int head, int child, int label) nogil:
+    if gold.labels[child] == -1:
+        return True
+    elif label == -1:
+        return True
+    elif gold.labels[child] == label:
+        return True
+    else:
+        return False
+
+
+cdef bint _is_gold_root(const GoldParseC* gold, int word) nogil:
+    return gold.labels[word] == -1 or gold.heads[word] == word
+ 
+
+cdef class Shift:
+    @staticmethod
+    cdef bint is_valid(StateClass st, int label) nogil:
+        return st.buffer_length() >= 2 and not st.shifted[st.B(0)] and not st.B_(0).sent_start
+
+    @staticmethod
+    cdef int transition(StateClass st, int label) nogil:
+        st.push()
+        st.fast_forward()
+
+    @staticmethod
+    cdef int cost(StateClass st, const GoldParseC* gold, int label) nogil:
+        return Shift.move_cost(st, gold) + Shift.label_cost(st, gold, label)
+
+    @staticmethod
+    cdef inline int move_cost(StateClass s, const GoldParseC* gold) nogil:
+        return push_cost(s, gold, s.B(0))
+
+    @staticmethod
+    cdef inline int label_cost(StateClass s, const GoldParseC* gold, int label) nogil:
+        return 0
+
+
+cdef class Reduce:
+    @staticmethod
+    cdef bint is_valid(StateClass st, int label) nogil:
+        return st.stack_depth() >= 2
+
+    @staticmethod
+    cdef int transition(StateClass st, int label) nogil:
+        if st.has_head(st.S(0)):
+            st.pop()
+        else:
+            st.unshift()
+        st.fast_forward()
+
+    @staticmethod
+    cdef int cost(StateClass s, const GoldParseC* gold, int label) nogil:
+        return Reduce.move_cost(s, gold) + Reduce.label_cost(s, gold, label)
+
+    @staticmethod
+    cdef inline int move_cost(StateClass st, const GoldParseC* gold) nogil:
+        return pop_cost(st, gold, st.S(0))
+
+    @staticmethod
+    cdef inline int label_cost(StateClass s, const GoldParseC* gold, int label) nogil:
+        return 0
+
+
+cdef class LeftArc:
+    @staticmethod
+    cdef bint is_valid(StateClass st, int label) nogil:
+        return not st.B_(0).sent_start
+
+    @staticmethod
+    cdef int transition(StateClass st, int label) nogil:
+        st.add_arc(st.B(0), st.S(0), label)
+        st.pop()
+        st.fast_forward()
+
+    @staticmethod
+    cdef int cost(StateClass s, const GoldParseC* gold, int label) nogil:
+        return LeftArc.move_cost(s, gold) + LeftArc.label_cost(s, gold, label)
+
+    @staticmethod
+    cdef inline int move_cost(StateClass s, const GoldParseC* gold) nogil:
+        cdef int cost = 0
+        if arc_is_gold(gold, s.B(0), s.S(0)):
+            return 0
+        else:
+            # Account for deps we might lose between S0 and stack
+            if not s.has_head(s.S(0)):
+                for i in range(1, s.stack_depth()):
+                    cost += gold.heads[s.S(i)] == s.S(0)
+                    cost += gold.heads[s.S(0)] == s.S(i)
+            return pop_cost(s, gold, s.S(0)) + arc_cost(s, gold, s.B(0), s.S(0))
+
+    @staticmethod
+    cdef inline int label_cost(StateClass s, const GoldParseC* gold, int label) nogil:
+        return arc_is_gold(gold, s.B(0), s.S(0)) and not label_is_gold(gold, s.B(0), s.S(0), label)
+
+
+cdef class RightArc:
+    @staticmethod
+    cdef bint is_valid(StateClass st, int label) nogil:
+        return not st.B_(0).sent_start
+
+    @staticmethod
+    cdef int transition(StateClass st, int label) nogil:
+        st.add_arc(st.S(0), st.B(0), label)
+        st.push()
+        st.fast_forward()
+
+    @staticmethod
+    cdef inline int cost(StateClass s, const GoldParseC* gold, int label) nogil:
+        return RightArc.move_cost(s, gold) + RightArc.label_cost(s, gold, label)
+
+    @staticmethod
+    cdef inline int move_cost(StateClass s, const GoldParseC* gold) nogil:
+        if arc_is_gold(gold, s.S(0), s.B(0)):
+            return 0
+        elif s.shifted[s.B(0)]:
+            return push_cost(s, gold, s.B(0))
+        else:
+            return push_cost(s, gold, s.B(0)) + arc_cost(s, gold, s.S(0), s.B(0))
+
+    @staticmethod
+    cdef int label_cost(StateClass s, const GoldParseC* gold, int label) nogil:
+        return arc_is_gold(gold, s.S(0), s.B(0)) and not label_is_gold(gold, s.S(0), s.B(0), label)
+
+
+cdef class Break:
+    @staticmethod
+    cdef bint is_valid(StateClass st, int label) nogil:
+        cdef int i
+        if not USE_BREAK:
+            return False
+        elif st.at_break():
+            return False
+        elif st.B(0) == 0:
+            return False
+        elif st.stack_depth() < 1:
+            return False
+        elif (st.S(0) + 1) != st.B(0):
+            # Must break at the token boundary
+            return False
+        else:
+            return True
+
+    @staticmethod
+    cdef int transition(StateClass st, int label) nogil:
+        st.set_break(st.B(0))
+        st.fast_forward()
+
+    @staticmethod
+    cdef int cost(StateClass s, const GoldParseC* gold, int label) nogil:
+        return Break.move_cost(s, gold) + Break.label_cost(s, gold, label)
+
+    @staticmethod
+    cdef inline int move_cost(StateClass s, const GoldParseC* gold) nogil:
+        cdef int cost = 0
+        cdef int i, j, S_i, B_i
+        for i in range(s.stack_depth()):
+            S_i = s.S(i)
+            for j in range(s.buffer_length()):
+                B_i = s.B(j)
+                cost += gold.heads[S_i] == B_i
+                cost += gold.heads[B_i] == S_i
+        # Check for sentence boundary --- if it's here, we can't have any deps
+        # between stack and buffer, so rest of action is irrelevant.
+        s0_root = _get_root(s.S(0), gold)
+        b0_root = _get_root(s.B(0), gold)
+        if s0_root != b0_root or s0_root == -1 or b0_root == -1:
+            return cost
+        else:
+            return cost + 1
+    
+    @staticmethod
+    cdef inline int label_cost(StateClass s, const GoldParseC* gold, int label) nogil:
+        return 0
+
+cdef int _get_root(int word, const GoldParseC* gold) nogil:
+    while gold.heads[word] != word and gold.labels[word] != -1 and word >= 0:
+        word = gold.heads[word]
+    if gold.labels[word] == -1:
+        return -1
+    else:
+        return word
+        
+
+cdef class ArcEager(TransitionSystem):
+    @classmethod
+    def get_labels(cls, gold_parses):
+        move_labels = {SHIFT: {'': True}, REDUCE: {'': True}, RIGHT: {'ROOT': True},
+                       LEFT: {'ROOT': True}, BREAK: {'ROOT': True}}
+        for raw_text, sents in gold_parses:
+            for (ids, words, tags, heads, labels, iob), ctnts in sents:
+                for child, head, label in zip(ids, heads, labels):
+                    if label.upper() == 'ROOT':
+                        label = 'ROOT'
+                    if label != 'ROOT':
+                        if head < child:
+                            move_labels[RIGHT][label] = True
+                        elif head > child:
+                            move_labels[LEFT][label] = True
+        return move_labels
+
+    cdef int preprocess_gold(self, GoldParse gold) except -1:
+        for i in range(gold.length):
+            if gold.heads[i] is None: # Missing values
+                gold.c.heads[i] = i
+                gold.c.labels[i] = -1
+            else:
+                label = gold.labels[i]
+                if label.upper() == 'ROOT':
+                    label = 'ROOT'
+                gold.c.heads[i] = gold.heads[i]
+                gold.c.labels[i] = self.strings[label]
+        for end, brackets in gold.brackets.items():
+            for start, label_strs in brackets.items():
+                gold.c.brackets[start][end] = 1
+                for label_str in label_strs:
+                    # Add the encoded label to the set
+                    gold.brackets[end][start].add(self.strings[label_str])
+
+    cdef Transition lookup_transition(self, object name) except *:
+        if '-' in name:
+            move_str, label_str = name.split('-', 1)
+            label = self.label_ids[label_str]
+        else:
+            label = 0
+        move = MOVE_NAMES.index(move_str)
+        for i in range(self.n_moves):
+            if self.c[i].move == move and self.c[i].label == label:
+                return self.c[i]
+
+    def move_name(self, int move, int label):
+        label_str = self.strings[label]
+        if label_str:
+            return MOVE_NAMES[move] + '-' + label_str
+        else:
+            return MOVE_NAMES[move]
+
+    cdef Transition init_transition(self, int clas, int move, int label) except *:
+        # TODO: Apparent Cython bug here when we try to use the Transition()
+        # constructor with the function pointers
+        cdef Transition t
+        t.score = 0
+        t.clas = clas
+        t.move = move
+        t.label = label
+        if move == SHIFT:
+            t.is_valid = Shift.is_valid
+            t.do = Shift.transition
+            t.get_cost = Shift.cost
+        elif move == REDUCE:
+            t.is_valid = Reduce.is_valid
+            t.do = Reduce.transition
+            t.get_cost = Reduce.cost
+        elif move == LEFT:
+            t.is_valid = LeftArc.is_valid
+            t.do = LeftArc.transition
+            t.get_cost = LeftArc.cost
+        elif move == RIGHT:
+            t.is_valid = RightArc.is_valid
+            t.do = RightArc.transition
+            t.get_cost = RightArc.cost
+        elif move == BREAK:
+            t.is_valid = Break.is_valid
+            t.do = Break.transition
+            t.get_cost = Break.cost
+        else:
+            raise Exception(move)
+        return t
+
+    cdef int initialize_state(self, StateClass st) except -1:
+        # Ensure sent_start is set to 0 throughout
+        for i in range(st.length):
+            st._sent[i].sent_start = False
+            st._sent[i].l_edge = i
+            st._sent[i].r_edge = i
+        st.fast_forward()
+
+    cdef int finalize_state(self, StateClass st) except -1:
+        cdef int root_label = self.strings['ROOT']
+        for i in range(st.length):
+            if st._sent[i].head == 0 and st._sent[i].dep == 0:
+                st._sent[i].dep = root_label
+            # If we're not using the Break transition, we segment via root-labelled
+            # arcs between the root words.
+            elif USE_ROOT_ARC_SEGMENT and st._sent[i].dep == root_label:
+                st._sent[i].head = 0
+
+    cdef int set_valid(self, bint* output, StateClass stcls) except -1:
+        cdef bint[N_MOVES] is_valid
+        is_valid[SHIFT] = Shift.is_valid(stcls, -1)
+        is_valid[REDUCE] = Reduce.is_valid(stcls, -1)
+        is_valid[LEFT] = LeftArc.is_valid(stcls, -1)
+        is_valid[RIGHT] = RightArc.is_valid(stcls, -1)
+        is_valid[BREAK] = Break.is_valid(stcls, -1)
+        cdef int i
+        n_valid = 0
+        for i in range(self.n_moves):
+            output[i] = is_valid[self.c[i].move]
+            n_valid += output[i]
+        assert n_valid >= 1
+
+    cdef int set_costs(self, int* output, StateClass stcls, GoldParse gold) except -1:
+        cdef int i, move, label
+        cdef label_cost_func_t[N_MOVES] label_cost_funcs
+        cdef move_cost_func_t[N_MOVES] move_cost_funcs
+        cdef int[N_MOVES] move_costs
+        for i in range(N_MOVES):
+            move_costs[i] = -1
+        move_cost_funcs[SHIFT] = Shift.move_cost
+        move_cost_funcs[REDUCE] = Reduce.move_cost
+        move_cost_funcs[LEFT] = LeftArc.move_cost
+        move_cost_funcs[RIGHT] = RightArc.move_cost
+        move_cost_funcs[BREAK] = Break.move_cost
+
+        label_cost_funcs[SHIFT] = Shift.label_cost
+        label_cost_funcs[REDUCE] = Reduce.label_cost
+        label_cost_funcs[LEFT] = LeftArc.label_cost
+        label_cost_funcs[RIGHT] = RightArc.label_cost
+        label_cost_funcs[BREAK] = Break.label_cost
+
+        cdef int* labels = gold.c.labels
+        cdef int* heads = gold.c.heads
+
+        n_gold = 0
+        for i in range(self.n_moves):
+            if self.c[i].is_valid(stcls, self.c[i].label):
+                move = self.c[i].move
+                label = self.c[i].label
+                if move_costs[move] == -1:
+                    move_costs[move] = move_cost_funcs[move](stcls, &gold.c)
+                output[i] = move_costs[move] + label_cost_funcs[move](stcls, &gold.c, label)
+                n_gold += output[i] == 0
+            else:
+                output[i] = 9000
+        assert n_gold >= 1
+
+    cdef Transition best_valid(self, const weight_t* scores, StateClass stcls) except *:
+        cdef bint[N_MOVES] is_valid
+        is_valid[SHIFT] = Shift.is_valid(stcls, -1)
+        is_valid[REDUCE] = Reduce.is_valid(stcls, -1)
+        is_valid[LEFT] = LeftArc.is_valid(stcls, -1)
+        is_valid[RIGHT] = RightArc.is_valid(stcls, -1)
+        is_valid[BREAK] = Break.is_valid(stcls, -1)
+        cdef Transition best
+        cdef weight_t score = MIN_SCORE
+        cdef int i
+        for i in range(self.n_moves):
+            if scores[i] > score and is_valid[self.c[i].move]:
+                best = self.c[i]
+                score = scores[i]
+        assert best.clas < self.n_moves
+        assert score > MIN_SCORE, (stcls.stack_depth(), stcls.buffer_length(), stcls.is_final(), stcls._b_i, stcls.length)
+        return best

spacy/syntax/ner.pyx

@@ -128,27 +128,6 @@ cdef class BiluoPushDown(TransitionSystem):
             raise Exception(move)
         return t
 
-    cdef Transition best_valid(self, const weight_t* scores, StateClass stcls) except *:
-        cdef int best = -1
-        cdef weight_t score = -90000
-        cdef const Transition* m
-        cdef int i
-        for i in range(self.n_moves):
-            m = &self.c[i]
-            if m.is_valid(stcls, m.label) and scores[i] > score:
-                best = i
-                score = scores[i]
-        assert best >= 0
-        cdef Transition t = self.c[best]
-        t.score = score
-        return t
-
-    cdef int set_valid(self, bint* output, StateClass stcls) except -1:
-        cdef int i
-        for i in range(self.n_moves):
-            m = &self.c[i]
-            output[i] = m.is_valid(stcls, m.label)
-
 
 cdef class Missing:
     @staticmethod

spacy/syntax/parser.pxd

@@ -8,9 +8,6 @@ from ..tokens cimport Tokens, TokenC
 
 
 cdef class Parser:
-    cdef readonly object cfg
-    cdef readonly Model model
-    cdef readonly TransitionSystem moves
-
-    cdef int _greedy_parse(self, Tokens tokens) except -1
-    cdef int _beam_parse(self, Tokens tokens) except -1
+    cdef public object cfg
+    cdef public Model model
+    cdef public TransitionSystem moves

spacy/syntax/parser.pyx

@@ -19,17 +19,10 @@ from cymem.cymem cimport Pool, Address
 from murmurhash.mrmr cimport hash64
 from thinc.typedefs cimport weight_t, class_t, feat_t, atom_t, hash_t
 
-
 from util import Config
 
-from thinc.features cimport Extractor
-from thinc.features cimport Feature
-from thinc.features cimport count_feats
+from thinc.api cimport Example
 
-from thinc.learner cimport LinearModel
-
-from thinc.search cimport Beam
-from thinc.search cimport MaxViolation
 
 from ..tokens cimport Tokens, TokenC
 from ..strings cimport StringStore
@@ -58,6 +51,8 @@ def get_templates(name):
         return pf.ner
     elif name == 'debug':
         return pf.unigrams
+    elif name.startswith('embed'):
+        return (pf.words, pf.tags, pf.labels)
     else:
         return (pf.unigrams + pf.s0_n0 + pf.s1_n0 + pf.s1_s0 + pf.s0_n1 + pf.n0_n1 + \
                 pf.tree_shape + pf.trigrams)
@@ -68,39 +63,103 @@ cdef class Parser:
         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)
+        self.model = Model(self.moves.n_moves, self.cfg.templates, model_dir)
 
     def __call__(self, Tokens tokens):
-        if self.cfg.get('beam_width', 1) < 1:
-            self._greedy_parse(tokens)
-        else:
-            self._beam_parse(tokens)
+        cdef StateClass stcls = StateClass.init(tokens.data, tokens.length)
+        self.moves.initialize_state(stcls)
+
+        cdef Example eg = Example(self.model.n_classes, CONTEXT_SIZE,
+                                  self.model.n_feats, self.model.n_feats)
+        eg.scores[0] = 10
+        assert eg.c.scores[0] == 10
+        while not stcls.is_final():
+            memset(eg.c.scores, 0, eg.c.nr_class * sizeof(weight_t))
+        
+            self.moves.set_valid(<bint*>eg.c.is_valid, stcls)
+            fill_context(eg.c.atoms, stcls)
+            
+            self.model.predict(eg)
+
+            self.moves.c[eg.c.guess].do(stcls, self.moves.c[eg.c.guess].label)
+        self.moves.finalize_state(stcls)
+        tokens.set_parse(stcls._sent)
 
     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 StateClass stcls = StateClass.init(tokens.data, tokens.length)
         self.moves.initialize_state(stcls)
-        cdef Transition guess
+        cdef Example eg = Example(self.model.n_classes, CONTEXT_SIZE,
+                                  self.model.n_feats, self.model.n_feats)
+        cdef weight_t loss = 0
         words = [w.orth_ for w in tokens]
+        cdef Transition G
         while not stcls.is_final():
-            fill_context(context, stcls)
-            scores = self.model.score(context)
-            guess = self.moves.best_valid(scores, stcls)
-            #print self.moves.move_name(guess.move, guess.label), stcls.print_state(words)
-            guess.do(stcls, guess.label)
-            assert stcls._s_i >= 0
-        self.moves.finalize_state(stcls)
-        tokens.set_parse(stcls._sent)
+            memset(eg.c.scores, 0, eg.c.nr_class * sizeof(weight_t))
+        
+            self.moves.set_costs(<bint*>eg.c.is_valid, eg.c.costs, stcls, gold)
+            
+            fill_context(eg.c.atoms, stcls)
+
+            self.model.train(eg)
+
+            G = self.moves.c[eg.c.guess]
+
+            #if eg.c.cost != 0:
+            #    print self.moves.move_name(G.move, G.label), stcls.print_state(words)
+            self.moves.c[eg.c.guess].do(stcls, self.moves.c[eg.c.guess].label)
+            loss += eg.c.loss
+        return loss
+
+
+
+# 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 = <StateClass>_dest
+    src = <StateClass>_src
+    moves = <const Transition*>_moves
+    dest.clone(src)
+    moves[clas].do(dest, moves[clas].label)
+
+
+cdef void* _init_state(Pool mem, int length, void* tokens) except NULL:
+    cdef StateClass st = StateClass.init(<const TokenC*>tokens, length)
+    st.fast_forward()
+    Py_INCREF(st)
+    return <void*>st
+
+
+cdef int _check_final_state(void* _state, void* extra_args) except -1:
+    return (<StateClass>_state).is_final()
+
+
+def _cleanup(Beam beam):
+    for i in range(beam.width):
+        Py_XDECREF(<PyObject*>beam._states[i].content)
+        Py_XDECREF(<PyObject*>beam._parents[i].content)
+
+cdef hash_t _hash_state(void* _state, void* _) except 0:
+    return <hash_t>_state
+    
+    #state = <const State*>_state
+    #cdef atom_t[10] rep
+
+    #rep[0] = state.stack[0] if state.stack_len >= 1 else 0
+    #rep[1] = state.stack[-1] if state.stack_len >= 2 else 0
+    #rep[2] = state.stack[-2] if state.stack_len >= 3 else 0
+    #rep[3] = state.i
+    #rep[4] = state.sent[state.stack[0]].l_kids if state.stack_len >= 1 else 0
+    #rep[5] = state.sent[state.stack[0]].r_kids if state.stack_len >= 1 else 0
+    #rep[6] = state.sent[state.stack[0]].dep if state.stack_len >= 1 else 0
+    #rep[7] = state.sent[state.stack[-1]].dep if state.stack_len >= 2 else 0
+    #if get_left(state, get_n0(state), 1) != NULL:
+    #    rep[8] = get_left(state, get_n0(state), 1).dep 
+    #else:
+    #    rep[8] = 0
+    #rep[9] = state.sent[state.i].l_kids
+    #return hash64(rep, sizeof(atom_t) * 10, 0)
+
 
     cdef int _beam_parse(self, Tokens tokens) except -1:
         cdef Beam beam = Beam(self.moves.n_moves, self.cfg.beam_width)
@@ -114,30 +173,6 @@ cdef class Parser:
         tokens.set_parse(state._sent)
         _cleanup(beam)
 
-    def _greedy_train(self, Tokens tokens, GoldParse gold):
-        cdef Pool mem = Pool()
-        cdef StateClass stcls = StateClass.init(tokens.data, tokens.length)
-        self.moves.initialize_state(stcls)
-
-        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
-        words = [w.orth_ for w in tokens]
-        history = []
-        while not stcls.is_final():
-            fill_context(context, stcls)
-            scores = self.model.score(context)
-            guess = self.moves.best_valid(scores, stcls)
-            best = self.moves.best_gold(scores, stcls, gold)
-            cost = guess.get_cost(stcls, &gold.c, guess.label)
-            self.model.update(context, guess.clas, best.clas, cost)
-            guess.do(stcls, guess.label)
-            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)
@@ -200,50 +235,4 @@ cdef class Parser:
             count_feats(counts[clas], feats, n_feats, inc)
             self.moves.c[clas].do(stcls, self.moves.c[clas].label)
 
-
-# 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 = <StateClass>_dest
-    src = <StateClass>_src
-    moves = <const Transition*>_moves
-    dest.clone(src)
-    moves[clas].do(dest, moves[clas].label)
-
-
-cdef void* _init_state(Pool mem, int length, void* tokens) except NULL:
-    cdef StateClass st = StateClass.init(<const TokenC*>tokens, length)
-    st.fast_forward()
-    Py_INCREF(st)
-    return <void*>st
-
-
-cdef int _check_final_state(void* _state, void* extra_args) except -1:
-    return (<StateClass>_state).is_final()
-
-
-def _cleanup(Beam beam):
-    for i in range(beam.width):
-        Py_XDECREF(<PyObject*>beam._states[i].content)
-        Py_XDECREF(<PyObject*>beam._parents[i].content)
-
-cdef hash_t _hash_state(void* _state, void* _) except 0:
-    return <hash_t>_state
-    
-    #state = <const State*>_state
-    #cdef atom_t[10] rep
-
-    #rep[0] = state.stack[0] if state.stack_len >= 1 else 0
-    #rep[1] = state.stack[-1] if state.stack_len >= 2 else 0
-    #rep[2] = state.stack[-2] if state.stack_len >= 3 else 0
-    #rep[3] = state.i
-    #rep[4] = state.sent[state.stack[0]].l_kids if state.stack_len >= 1 else 0
-    #rep[5] = state.sent[state.stack[0]].r_kids if state.stack_len >= 1 else 0
-    #rep[6] = state.sent[state.stack[0]].dep if state.stack_len >= 1 else 0
-    #rep[7] = state.sent[state.stack[-1]].dep if state.stack_len >= 2 else 0
-    #if get_left(state, get_n0(state), 1) != NULL:
-    #    rep[8] = get_left(state, get_n0(state), 1).dep 
-    #else:
-    #    rep[8] = 0
-    #rep[9] = state.sent[state.i].l_kids
-    #return hash64(rep, sizeof(atom_t) * 10, 0)
+"""

spacy/syntax/transition_system.pxd

@@ -46,9 +46,5 @@ cdef class TransitionSystem:
 
     cdef int set_valid(self, bint* output, StateClass state) except -1
     
-    cdef int set_costs(self, int* output, StateClass state, GoldParse gold) except -1
-
-    cdef Transition best_valid(self, const weight_t* scores, StateClass stcls) except *
-
-    cdef Transition best_gold(self, const weight_t* scores, StateClass state,
-                              GoldParse gold) except *
+    cdef int set_costs(self, bint* is_valid, int* costs,
+                       StateClass state, GoldParse gold) except -1

spacy/syntax/transition_system.pyx

@@ -43,30 +43,17 @@ cdef class TransitionSystem:
     cdef Transition init_transition(self, int clas, int move, int label) except *:
         raise NotImplementedError
 
-    cdef Transition best_valid(self, const weight_t* scores, StateClass s) except *:
-        raise NotImplementedError
-    
-    cdef int set_valid(self, bint* output, StateClass state) except -1:
-        raise NotImplementedError
-
-    cdef int set_costs(self, int* output, StateClass stcls, GoldParse gold) except -1:
+    cdef int set_valid(self, bint* is_valid, StateClass stcls) except -1:
         cdef int i
         for i in range(self.n_moves):
-            if self.c[i].is_valid(stcls, self.c[i].label):
-                output[i] = self.c[i].get_cost(stcls, &gold.c, self.c[i].label)
+            is_valid[i] = self.c[i].is_valid(stcls, self.c[i].label)
+
+    cdef int set_costs(self, bint* is_valid, int* costs,
+                       StateClass stcls, GoldParse gold) except -1:
+        cdef int i
+        self.set_valid(is_valid, stcls)
+        for i in range(self.n_moves):
+            if is_valid[i]:
+                costs[i] = self.c[i].get_cost(stcls, &gold.c, self.c[i].label)
             else:
-                output[i] = 9000
-
-    cdef Transition best_gold(self, const weight_t* scores, StateClass stcls,
-                              GoldParse gold) except *:
-        cdef Transition best
-        cdef weight_t score = MIN_SCORE
-        cdef int i
-        for i in range(self.n_moves):
-            if self.c[i].is_valid(stcls, self.c[i].label):
-                cost = self.c[i].get_cost(stcls, &gold.c, self.c[i].label)
-                if scores[i] > score and cost == 0:
-                    best = self.c[i]
-                    score = scores[i]
-        assert score > MIN_SCORE
-        return best
+                costs[i] = 9000

spacy/tokens.pyx

@@ -15,6 +15,7 @@ from .parts_of_speech cimport CONJ, PUNCT
 from .lexeme cimport check_flag
 from .spans import Span
 from .structs cimport UniStr
+from .senses import STRINGS as SENSE_STRINGS
 
 from unidecode import unidecode
 # Compiler crashes on memory view coercion without this. Should report bug.
@@ -92,7 +93,7 @@ cdef class Tokens:
         else:
             size = 5
         self.mem = Pool()
-        # Guarantee self.lex[i-x], for any i >= 0 and x < padding is in bounds
+        # Guarantee self.data[i-x], for any i >= 0 and x < padding is in bounds
         # However, we need to remember the true starting places, so that we can
         # realloc.
         data_start = <TokenC*>self.mem.alloc(size + (PADDING*2), sizeof(TokenC))
@@ -462,6 +463,10 @@ cdef class Token:
         def __get__(self):
             return self.c.dep
 
+    property sense:
+        def __get__(self):
+            return self.c.sense
+
     property repvec:
         def __get__(self):
             cdef int length = self.vocab.repvec_length
@@ -646,6 +651,11 @@ cdef class Token:
         def __get__(self):
             return self.vocab.strings[self.c.dep]
 
+    property sense_:
+        def __get__(self):
+            return SENSE_STRINGS[self.c.sense]
+
+
 
 _pos_id_to_string = {id_: string for string, id_ in UNIV_POS_NAMES.items()}
 

spacy/wsd/supersense_tagger.pyx

@@ -0,0 +1,431 @@
+from libc.string cimport memcpy
+from libc.math cimport exp
+
+from cymem.cymem cimport Pool
+
+from thinc.learner cimport LinearModel
+from thinc.features cimport Extractor, Feature
+
+from thinc.typedefs cimport atom_t, weight_t, feat_t
+cimport cython
+
+
+from ..typedefs cimport flags_t
+from ..structs cimport TokenC
+from ..strings cimport StringStore
+from ..tokens cimport Tokens
+from .supersenses cimport N_SENSES, encode_supersense_strs
+from .supersenses cimport NO_SENSE, N_Tops, J_all, J_pert, A_all, J_ppl, V_body
+from ..gold cimport GoldParse
+from ..parts_of_speech cimport NOUN, VERB, ADV, ADJ, N_UNIV_TAGS
+
+from .. cimport parts_of_speech
+
+from os import path
+import json
+
+
+cdef enum:
+    P2W
+    P2p
+    P2c
+    P2c6
+    P2c4
+
+    P1W
+    P1p
+    P1c
+    P1c6
+    P1c4
+
+    N0W
+    N0p
+    N0c
+    N0c6
+    N0c4
+    
+    N1W
+    N1p
+    N1c
+    N1c6
+    N1c4
+    
+    N2W
+    N2p
+    N2c
+    N2c6
+    N2c4
+ 
+    Hw
+    Hp
+    Hc
+    Hc6
+    Hc4
+    
+    N3W
+    P3W
+
+    P1s
+    P2s
+
+   
+    CONTEXT_SIZE
+
+
+unigrams = (
+    (Hw,),
+    (Hp,),
+    (Hw, Hp),
+    (Hc, Hp),
+    (Hc6, Hp),
+    (Hc4, Hp),
+    (Hc,),
+
+    (P2W,),
+    (P2p,),
+    (P2W, P2p),
+    (P2c, P2p),
+    (P2c6, P2p),
+    (P2c4, P2p),
+    (P2c,),
+
+    (P1W,),
+    (P1p,),
+    (P1W, P1p),
+    (P1c, P1p),
+    (P1c6, P1p),
+    (P1c4, P1p),
+    (P1c,),
+    (P1W,),
+    (P1p,),
+    (P1W, P1p),
+    (P1c, P1p),
+    (P1c6, P1p),
+    (P1c4, P1p),
+    (P1c,),
+    
+    (N0p,),
+    (N0c, N0p),
+    (N0c6, N0p),
+    (N0c4, N0p),
+    (N0c,),
+    (N0p,),
+    (N0c, N0p),
+    (N0c6, N0p),
+    (N0c4, N0p),
+    (N0c,),
+
+    (N1p,),
+    (N1W, N1p),
+    (N1c, N1p),
+    (N1c6, N1p),
+    (N1c4, N1p),
+    (N1c,),
+    (N1W,),
+    (N1p,),
+    (N1W, N1p),
+    (N1c, N1p),
+    (N1c6, N1p),
+    (N1c4, N1p),
+    (N1c,),
+
+    (N2p,),
+    (N2W, N2p),
+    (N2c, N2p),
+    (N2c6, N2p),
+    (N2c4, N2p),
+    (N2c,),
+    (N2W,),
+    (N2p,),
+    (N2W, N2p),
+    (N2c, N2p),
+    (N2c6, N2p),
+    (N2c4, N2p),
+    (N2c,),
+
+    (P1s,),
+    (P2s,),
+    (P1s, P2s,),
+    (P1s, N0p),
+    (P1s, P2s, N0c),
+
+    (N3W,),
+    (P3W,),
+)
+
+
+bigrams = (
+    (P2p, P1p),
+    (P2W, N0p),
+    (P2c, P1p),
+    (P1c, N0p),
+    (P1c6, N0p),
+
+    (N0p, N1p,),
+    (P2W, P1W),
+    (P1W, N1W),
+    (N1W, N2W),
+)
+
+
+trigrams = (
+    (P1p, N0p, N1p),
+    (P2p, P1p,),
+    (P2c4, P1c4, N0c4),
+    
+    (P1p, N0p, N1p),
+    (P1p, N0p,),
+    (P1c4, N0c4, N1c4),
+
+    (N0p, N1p, N2p),
+    (N0p, N1p,),
+    (N0c4, N1c4, N2c4),
+    (P1W, N0p, N0W),
+)
+
+
+cdef int fill_token(atom_t* ctxt, const TokenC* token) except -1:
+    ctxt[0] = token.lemma
+    ctxt[1] = token.tag
+    ctxt[2] = token.lex.cluster
+    ctxt[3] = token.lex.cluster & 15
+    ctxt[4] = token.lex.cluster & 63
+
+
+cdef int fill_context(atom_t* ctxt, const TokenC* token) except -1:
+    # NB: we have padding to keep us safe here
+    # See tokens.pyx
+    fill_token(&ctxt[P2W], token - 2)
+    fill_token(&ctxt[P1W], token - 1)
+
+    fill_token(&ctxt[N0W], token)
+    ctxt[N0W] = 0 # Important! Don't condition on this
+
+    fill_token(&ctxt[N1W], token + 1)
+    fill_token(&ctxt[N2W], token + 2)
+    fill_token(&ctxt[Hw], token + token.head)
+    ctxt[P1s] = (token - 1).sense
+    ctxt[P2s] = (token - 2).sense
+    ctxt[N3W] = (token + 3).lemma
+    ctxt[P3W] = (token - 3).lemma
+
+
+cdef class FeatureVector:
+    cdef Pool mem
+    cdef Feature* c
+    cdef list extractors
+    cdef int length
+    cdef int _max_length
+
+    def __init__(self, length=100):
+        self.mem = Pool()
+        self.c = <Feature*>self.mem.alloc(length, sizeof(Feature))
+        self.length = 0
+        self._max_length = length
+
+    def __len__(self):
+        return self.length
+
+    cpdef int add(self, feat_t key, weight_t value) except -1:
+        if self.length == self._max_length:
+            self._max_length *= 2
+            self.c = <Feature*>self.mem.realloc(self.c, self._max_length * sizeof(Feature))
+
+        self.c[self.length] = Feature(i=0, key=key, value=value)
+        self.length += 1
+
+    cdef int extend(self, const Feature* new_feats, int n_feats) except -1:
+        new_length = self.length + n_feats
+        if new_length >= self._max_length:
+            self._max_length = 2 * new_length
+            self.c = <Feature*>self.mem.realloc(self.c, new_length * sizeof(Feature))
+        memcpy(&self.c[self.length], new_feats, n_feats * sizeof(Feature))
+        self.length += n_feats
+
+    def clear(self):
+        self.length = 0
+
+
+cdef class SenseTagger:
+    cdef readonly StringStore strings
+    cdef readonly LinearModel model
+    cdef readonly Extractor extractor
+    cdef readonly model_dir
+    cdef readonly flags_t[<int>N_UNIV_TAGS] pos_senses
+    cdef dict tagdict
+
+    def __init__(self, StringStore strings, model_dir):
+        self.model_dir = model_dir
+        if path.exists(path.join(model_dir, 'wordnet', 'supersenses.json')):
+            self.tagdict = json.load(open(path.join(model_dir, 'wordnet', 'supersenses.json')))
+        else:
+            self.tagdict = {}
+
+        if model_dir is not None and path.isdir(model_dir):
+            model_dir = path.join(model_dir, 'wsd')
+
+        templates = unigrams + bigrams + trigrams
+        self.extractor = Extractor(templates)
+        self.model = LinearModel(N_SENSES, self.extractor.n_templ)
+
+        self.strings = strings
+        cdef flags_t all_senses = 0
+        cdef flags_t sense = 0
+        cdef flags_t one = 1
+        for sense in range(1, N_SENSES):
+            all_senses |= (one << sense)
+
+        self.pos_senses[<int>parts_of_speech.NO_TAG] = all_senses
+        self.pos_senses[<int>parts_of_speech.ADJ] = all_senses
+        self.pos_senses[<int>parts_of_speech.ADV] = all_senses
+        self.pos_senses[<int>parts_of_speech.ADP] = all_senses
+        self.pos_senses[<int>parts_of_speech.CONJ] = 0
+        self.pos_senses[<int>parts_of_speech.DET] = 0
+        self.pos_senses[<int>parts_of_speech.NUM] = 0
+        self.pos_senses[<int>parts_of_speech.PRON] = 0
+        self.pos_senses[<int>parts_of_speech.PRT] = all_senses
+        self.pos_senses[<int>parts_of_speech.X] = all_senses
+        self.pos_senses[<int>parts_of_speech.PUNCT] = 0
+        self.pos_senses[<int>parts_of_speech.EOL] = 0
+
+        for sense in range(N_Tops, V_body):
+            self.pos_senses[<int>parts_of_speech.NOUN] |= one << sense
+
+        self.pos_senses[<int>parts_of_speech.VERB] = 0
+        for sense in range(V_body, J_ppl):
+            self.pos_senses[<int>parts_of_speech.VERB] |= one << sense
+
+    def __call__(self, Tokens tokens):
+        cdef atom_t[CONTEXT_SIZE] local_context
+        cdef int i, guess, n_feats
+        cdef flags_t valid_senses = 0
+        cdef TokenC* token
+        cdef flags_t one = 1
+        cdef int n_doc_feats
+        cdef Pool mem = Pool()
+        feats = self.get_doc_feats(mem, tokens, &n_doc_feats)
+        for i in range(tokens.length):
+            token = &tokens.data[i]
+            valid_senses = token.lex.senses & self.pos_senses[<int>token.pos]
+            if valid_senses >= 2:
+                fill_context(local_context, token)
+                n_local_feats = self.extractor.set_feats(&feats[n_doc_feats],
+                                                         local_context)
+                scores = self.model.get_scores(feats, n_local_feats)
+                self.weight_scores_by_tagdict(<weight_t*><void*>scores, token, 0.0)
+                tokens.data[i].sense = self.best_in_set(scores, valid_senses)
+            else:
+                token.sense = NO_SENSE
+
+    def train(self, Tokens tokens):
+        cdef int i, j
+        cdef TokenC* token
+        cdef atom_t[CONTEXT_SIZE] context
+        cdef int n_doc_feats, n_local_feats
+        cdef feat_t f_key
+        cdef flags_t best_senses = 0
+        cdef int f_i
+        cdef int cost = 0
+
+        cdef Pool mem = Pool()
+        feats = self.get_doc_feats(mem, tokens, &n_doc_feats)
+        for i in range(tokens.length):
+            token = &tokens.data[i]
+            pos_senses = self.pos_senses[<int>token.pos]
+            lex_senses = token.lex.senses & pos_senses
+            if lex_senses >= 2:
+                fill_context(context, token)
+
+                n_local_feats = self.extractor.set_feats(&feats[n_doc_feats], context)
+                scores = self.model.get_scores(feats, n_doc_feats + n_local_feats)
+                guess  = self.best_in_set(scores, pos_senses)
+                best   = self.best_in_set(scores, lex_senses)
+                update = self._make_update(feats, n_doc_feats + n_local_feats,
+                                           guess, best)
+                self.model.update(update)
+                token.sense = best
+                cost += guess != best
+            else:
+                token.sense = 1
+        return cost
+
+    cdef dict _make_update(self, const Feature* feats, int n_feats, int guess, int best):
+        guess_counts = {}
+        gold_counts = {}
+        if guess != best:
+            for j in range(n_feats):
+                f_key = feats[j].key
+                f_i = feats[j].i
+                feat = (f_i, f_key)
+                gold_counts[feat]  = gold_counts.get(feat, 0) + 1.0
+                guess_counts[feat] = guess_counts.get(feat, 0) - 1.0
+        return {guess: guess_counts, best: gold_counts}
+
+    cdef Feature* get_doc_feats(self, Pool mem, Tokens tokens, int* n_feats) except NULL:
+        # Get features for the document
+        # Start with activation strengths for each supersense
+        n_feats[0] = N_SENSES
+        feats = <Feature*>mem.alloc(n_feats[0] + self.extractor.n_templ + 1,
+                                    sizeof(Feature))
+        cdef int i, ssense
+        for ssense in range(N_SENSES):
+            feats[ssense] = Feature(i=0, key=ssense, value=0)
+        cdef flags_t pos_senses
+        cdef flags_t one = 1
+        for i in range(tokens.length):
+            sense_probs = self.tagdict.get(tokens.data[i].lemma, {})
+            pos_senses = self.pos_senses[<int>tokens.data[i].pos]
+            for ssense_str, prob in sense_probs.items():
+                ssense = int(ssense_str + 1)
+                if pos_senses & (one << <flags_t>ssense):
+                    feats[ssense].value += prob
+        return feats
+
+    cdef int best_in_set(self, const weight_t* scores, flags_t senses) except -1:
+        cdef weight_t max_ = 0
+        cdef int argmax = -1
+        cdef flags_t i
+        cdef flags_t one = 1
+        for i in range(N_SENSES):
+            if (senses & (one << i)) and (argmax == -1 or scores[i] > max_):
+                max_ = scores[i]
+                argmax = i
+        assert argmax >= 0
+        return argmax
+
+    cdef int weight_scores_by_tagdict(self, weight_t* scores, const TokenC* token,
+                                      weight_t a) except -1:
+        lemma = self.strings[token.lemma]
+
+        # First softmax the scores
+        softmax(scores, N_SENSES)
+
+        probs = self.tagdict.get(lemma, {})
+        for i in range(1, N_SENSES):
+            prob = probs.get(unicode(i-1), 0)
+            scores[i] = (a * prob) + ((1 - a) * scores[i])
+
+    def end_training(self):
+        self.model.end_training()
+        self.model.dump(path.join(self.model_dir, 'model'), freq_thresh=0)
+
+
+@cython.cdivision(True)
+cdef void softmax(weight_t* scores, int n_classes) nogil:
+    cdef int i
+    cdef double total = 0
+    for i in range(N_SENSES):
+        total += exp(scores[i])
+    for i in range(N_SENSES):
+        scores[i] = <weight_t>(exp(scores[i]) / total)
+
+
+
+cdef list _set_bits(flags_t flags):
+    bits = []
+    cdef flags_t bit
+    cdef flags_t one = 1
+    for bit in range(N_SENSES):
+        if flags & (one << bit):
+            bits.append(bit)
+    return bits

spacy/wsd/supersenses.pxd

@@ -1,28 +1,20 @@
 # Enum of Wordnet supersenses
-cimport parts_of_speech
-from .typedefs cimport flags_t
+from ..typedefs cimport flags_t
+from .. cimport parts_of_speech
 
 cpdef enum:
-    A_behavior
-    A_body
-    A_feeling
-    A_mind
-    A_motion
-    A_perception
-    A_quantity
-    A_relation
-    A_social
-    A_spatial
-    A_substance
-    A_time
-    A_weather
+    NO_SENSE
+    J_all
+    J_pert
+    A_all
+    N_Tops
     N_act
     N_animal
     N_artifact
     N_attribute
     N_body
     N_cognition
-    N_communication
+    N_communication  
     N_event
     N_feeling
     N_food
@@ -56,7 +48,8 @@ cpdef enum:
     V_social
     V_stative
     V_weather
+    J_ppl
+    N_SENSES
 
 
-cdef flags_t[<int>parts_of_speech.N_UNIV_TAGS] POS_SENSES
-
+cdef flags_t encode_supersense_strs(sense_names) except 0

spacy/wsd/supersenses.pyx

@@ -0,0 +1,69 @@
+from __future__ import unicode_literals
+from .. cimport parts_of_speech
+
+
+lexnames_str = """
+-1      NO_SENSE       -1
+00      J_all 3
+01      A_pert        3 
+02      A_all 4
+03      N_Tops       1  
+04      N_act        1
+05      N_animal     1
+06      N_artifact   1
+07      N_attribute  1
+08      N_body       1
+09      N_cognition  1
+10      N_communication      1
+11      N_event      1
+12      N_feeling    1
+13      N_food       1
+14      N_group      1
+15      N_location   1
+16      N_motive     1
+17      N_object     1
+18      N_person     1
+19      N_phenomenon 1
+20      N_plant      1
+21      N_possession 1
+22      N_process    1
+23      N_quantity   1
+24      N_relation   1
+25      N_shape      1
+26      N_state      1
+27      N_substance  1
+28      N_time       1
+29      V_body       2
+30      V_change     2
+31      V_cognition  2
+32      V_communication      2
+33      V_competition        2
+34      V_consumption        2
+35      V_contact    2
+36      V_creation   2
+37      V_emotion    2
+38      V_motion     2
+39      V_perception 2
+40      V_possession 2
+41      V_social     2
+42      V_stative    2
+43      V_weather    2
+44      A_ppl 3
+""".strip()
+
+STRINGS = tuple(line.split()[1] for line in lexnames_str.split('\n'))
+
+IDS = dict((sense_str, i) for i, sense_str in enumerate(STRINGS))
+
+
+cdef flags_t encode_supersense_strs(sense_names) except 0:
+    cdef flags_t sense_bits = 0
+    if len(sense_names) == 0:
+        return sense_bits | (1 << NO_SENSE)
+    cdef flags_t sense_id = 0
+    for sense_str in sense_names:
+        sense_str = sense_str.replace('noun', 'N').replace('verb', 'V')
+        sense_str = sense_str.replace('adj', 'J').replace('adv', 'A')
+        sense_id = IDS[sense_str]
+        sense_bits |= (1 << sense_id)
+    return sense_bits