<commit_before>'''Train for CONLL 2017 UD treebank evaluation. Takes .conllu files, writes
.conllu format for development data, allowing the official scorer to be used.
'''
from __future__ import unicode_literals
import plac
import tqdm
from pathlib import Path
import re
import sys
import json

import spacy
import spacy.util
from ..tokens import Token, Doc
from ..gold import GoldParse
from ..util import compounding, minibatch_by_words
from ..syntax.nonproj import projectivize
from ..matcher import Matcher
from .. import displacy
from collections import defaultdict, Counter
from timeit import default_timer as timer

import itertools
import random
import numpy.random
import cytoolz

from . import conll17_ud_eval

from .. import lang
from .. import lang
from ..lang import zh
from ..lang import ja

lang.zh.Chinese.Defaults.use_jieba = False
lang.ja.Japanese.Defaults.use_janome = False

random.seed(0)
numpy.random.seed(0)

################
# Data reading #
################

space_re = re.compile('\s+')
def split_text(text):
    return [space_re.sub(' ', par.strip()) for par in text.split('\n\n')]

def read_data(nlp, conllu_file, text_file, raw_text=True, oracle_segments=False,
              max_doc_length=None, limit=None):
    '''Read the CONLLU format into (Doc, GoldParse) tuples. If raw_text=True,
    include Doc objects created using nlp.make_doc and then aligned against
    the gold-standard sequences. If oracle_segments=True, include Doc objects
    created from the gold-standard segments. At least one must be True.'''
    if not raw_text and not oracle_segments:
        raise ValueError("At least one of raw_text or oracle_segments must be True")
    paragraphs = split_text(text_file.read())
    conllu = read_conllu(conllu_file)
    # sd is spacy doc; cd is conllu doc
    # cs is conllu sent, ct is conllu token
    docs = []
    golds = []
    for doc_id, (text, cd) in enumerate(zip(paragraphs, conllu)):
        sent_annots = []
        for cs in cd:
            sent = defaultdict(list)
            for id_, word, lemma, pos, tag, morph, head, dep, _, space_after in cs:
                if '.' in id_:
                    continue
                if '-' in id_:
                    continue
                id_ = int(id_)-1
                head = int(head)-1 if head != '0' else id_
                sent['words'].append(word)
                sent['tags'].append(tag)
                sent['heads'].append(head)
                sent['deps'].append('ROOT' if dep == 'root' else dep)
                sent['spaces'].append(space_after == '_')
            sent['entities'] = ['-'] * len(sent['words'])
            sent['heads'], sent['deps'] = projectivize(sent['heads'],
                                                       sent['deps'])
            if oracle_segments:
                docs.append(Doc(nlp.vocab, words=sent['words'], spaces=sent['spaces']))
                golds.append(GoldParse(docs[-1], **sent))

            sent_annots.append(sent)
            if raw_text and max_doc_length and len(sent_annots) >= max_doc_length:
                doc, gold = _make_gold(nlp, None, sent_annots)
                sent_annots = []
                docs.append(doc)
                golds.append(gold)
                if limit and len(docs) >= limit:
                    return docs, golds

        if raw_text and sent_annots:
            doc, gold = _make_gold(nlp, None, sent_annots)
            docs.append(doc)
            golds.append(gold)
        if limit and len(docs) >= limit:
            return docs, golds
    return docs, golds

def read_conllu(file_):
    docs = []
    sent = []
    doc = []
    for line in file_:
        if line.startswith('# newdoc'):
            if doc:
                docs.append(doc)
            doc = []
        elif line.startswith('#'):
            continue
        elif not line.strip():
            if sent:
                doc.append(sent)
            sent = []
        else:
            sent.append(list(line.strip().split('\t')))
            if len(sent[-1]) != 10:
                print(repr(line))
                raise ValueError
    if sent:
        doc.append(sent)
    if doc:
        docs.append(doc)
    return docs

def _make_gold(nlp, text, sent_annots):
    # Flatten the conll annotations, and adjust the head indices
    flat = defaultdict(list)
    for sent in sent_annots:
        flat['heads'].extend(len(flat['words'])+head for head in sent['heads'])
        for field in ['words', 'tags', 'deps', 'entities', 'spaces']:
            flat[field].extend(sent[field])
    # Construct text if necessary
    assert len(flat['words']) == len(flat['spaces'])
    if text is None:
        text = ''.join(word+' '*space for word, space in zip(flat['words'], flat['spaces']))
    doc = nlp.make_doc(text)
    flat.pop('spaces')
    gold = GoldParse(doc, **flat)
    return doc, gold

#############################
# Data transforms for spaCy #
#############################

def golds_to_gold_tuples(docs, golds):
    '''Get out the annoying 'tuples' format used by begin_training, given the
    GoldParse objects.'''
    tuples = []
    for doc, gold in zip(docs, golds):
        text = doc.text
        ids, words, tags, heads, labels, iob = zip(*gold.orig_annot)
        sents = [((ids, words, tags, heads, labels, iob), [])]
        tuples.append((text, sents))
    return tuples

##############
# Evaluation #
##############

def evaluate(nlp, text_loc, gold_loc, sys_loc, limit=None):
    with text_loc.open('r', encoding='utf8') as text_file:
        texts = split_text(text_file.read())
        docs = list(nlp.pipe(texts))
    with sys_loc.open('w', encoding='utf8') as out_file:
        write_conllu(docs, out_file)
    with gold_loc.open('r', encoding='utf8') as gold_file:
        gold_ud = conll17_ud_eval.load_conllu(gold_file)
        with sys_loc.open('r', encoding='utf8') as sys_file:
            sys_ud = conll17_ud_eval.load_conllu(sys_file)
        scores = conll17_ud_eval.evaluate(gold_ud, sys_ud)
    return docs, scores

def write_conllu(docs, file_):
    merger = Matcher(docs[0].vocab)
    merger.add('SUBTOK', None, [{'DEP': 'subtok', 'op': '+'}])
    for i, doc in enumerate(docs):
        matches = merger(doc)
        spans = [doc[start:end+1] for _, start, end in matches]
        offsets = [(span.start_char, span.end_char) for span in spans]
        for start_char, end_char in offsets:
            doc.merge(start_char, end_char)
        file_.write("# newdoc id = {i}\n".format(i=i))
        for j, sent in enumerate(doc.sents):
            file_.write("# sent_id = {i}.{j}\n".format(i=i, j=j))
            file_.write("# text = {text}\n".format(text=sent.text))
            for k, token in enumerate(sent):
                file_.write(token._.get_conllu_lines(k) + '\n')
            file_.write('\n')

def print_progress(itn, losses, ud_scores):
    fields = {
        'dep_loss': losses.get('parser', 0.0),
        'tag_loss': losses.get('tagger', 0.0),
        'words': ud_scores['Words'].f1 * 100,
        'sents': ud_scores['Sentences'].f1 * 100,
        'tags': ud_scores['XPOS'].f1 * 100,
        'uas': ud_scores['UAS'].f1 * 100,
        'las': ud_scores['LAS'].f1 * 100,
    }
    header = ['Epoch', 'Loss', 'LAS', 'UAS', 'TAG', 'SENT', 'WORD']
    if itn == 0:
        print('\t'.join(header))
    tpl = '\t'.join((
        '{:d}',
        '{dep_loss:.1f}',
        '{las:.1f}',
        '{uas:.1f}',
        '{tags:.1f}',
        '{sents:.1f}',
        '{words:.1f}',
    ))
    print(tpl.format(itn, **fields))

#def get_sent_conllu(sent, sent_id):
#    lines = ["# sent_id = {sent_id}".format(sent_id=sent_id)]

def get_token_conllu(token, i):
    if token._.begins_fused:
        n = 1
        while token.nbor(n)._.inside_fused:
            n += 1
        id_ = '%d-%d' % (i, i+n)
        lines = [id_, token.text, '_', '_', '_', '_', '_', '_', '_', '_']
    else:
        lines = []
    if token.head.i == token.i:
        head = 0
    else:
        head = i + (token.head.i - token.i) + 1
    fields = [str(i+1), token.text, token.lemma_, token.pos_, token.tag_, '_',
              str(head), token.dep_.lower(), '_', '_']
    lines.append('\t'.join(fields))
    return '\n'.join(lines)

Token.set_extension('get_conllu_lines', method=get_token_conllu)
Token.set_extension('begins_fused', default=False)
Token.set_extension('inside_fused', default=False)

##################
# Initialization #
##################

def load_nlp(corpus, config):
    lang = corpus.split('_')[0]
    nlp = spacy.blank(lang)
    if config.vectors:
        nlp.vocab.from_disk(config.vectors / 'vocab')
    return nlp

def initialize_pipeline(nlp, docs, golds, config):
    nlp.add_pipe(nlp.create_pipe('parser'))
    if config.multitask_tag:
        nlp.parser.add_multitask_objective('tag')
    if config.multitask_sent:
        nlp.parser.add_multitask_objective('sent_start')
    nlp.add_pipe(nlp.create_pipe('tagger'))
    for gold in golds:
        for tag in gold.tags:
            if tag is not None:
                nlp.tagger.add_label(tag)
    return nlp.begin_training(lambda: golds_to_gold_tuples(docs, golds))

########################
# Command line helpers #
########################

class Config(object):
    def __init__(self, vectors=None, max_doc_length=10, multitask_tag=True,
            multitask_sent=True, nr_epoch=30, batch_size=1000, dropout=0.2):
        for key, value in locals().items():
            setattr(self, key, value)

    @classmethod
    def load(cls, loc):
        with Path(loc).open('r', encoding='utf8') as file_:
            cfg = json.load(file_)
        return cls(**cfg)

class Dataset(object):
    def __init__(self, path, section):
        self.path = path
        self.section = section
        self.conllu = None
        self.text = None
        for file_path in self.path.iterdir():
            name = file_path.parts[-1]
            if section in name and name.endswith('conllu'):
                self.conllu = file_path
            elif section in name and name.endswith('txt'):
                self.text = file_path
        if self.conllu is None:
            msg = "Could not find .txt file in {path} for {section}"
            raise IOError(msg.format(section=section, path=path))
        if self.text is None:
            msg = "Could not find .txt file in {path} for {section}"
        self.lang = self.conllu.parts[-1].split('-')[0].split('_')[0]

class TreebankPaths(object):
    def __init__(self, ud_path, treebank, **cfg):
        self.train = Dataset(ud_path / treebank, 'train')
        self.dev = Dataset(ud_path / treebank, 'dev')
        self.lang = self.train.lang

@plac.annotations(
    ud_dir=("Path to Universal Dependencies corpus", "positional", None, Path),
    corpus=("UD corpus to train and evaluate on, e.g. en, es_ancora, etc",
            "positional", None, str),
    parses_dir=("Directory to write the development parses", "positional", None, Path),
    config=("Path to json formatted config file", "positional"),
    limit=("Size limit", "option", "n", int)
)
def main(ud_dir, parses_dir, config, corpus, limit=0):
    config = Config.load(config)
    paths = TreebankPaths(ud_dir, corpus)
    if not (parses_dir / corpus).exists():
        (parses_dir / corpus).mkdir()
    print("Train and evaluate", corpus, "using lang", paths.lang)
    nlp = load_nlp(paths.lang, config)

    docs, golds = read_data(nlp, paths.train.conllu.open(), paths.train.text.open(),
                            max_doc_length=config.max_doc_length, limit=limit)

    optimizer = initialize_pipeline(nlp, docs, golds, config)

    batch_sizes = compounding(config.batch_size //10, config.batch_size, 1.001)
    for i in range(config.nr_epoch):
        docs = [nlp.make_doc(doc.text) for doc in docs]
        Xs = list(zip(docs, golds))
        random.shuffle(Xs)
        batches = minibatch_by_words(Xs, size=batch_sizes)
        losses = {}
        n_train_words = sum(len(doc) for doc in docs)
        with tqdm.tqdm(total=n_train_words, leave=False) as pbar:
            for batch in batches:
                batch_docs, batch_gold = zip(*batch)
                pbar.update(sum(len(doc) for doc in batch_docs))
                nlp.update(batch_docs, batch_gold, sgd=optimizer,
                           drop=config.dropout, losses=losses)

        out_path = parses_dir / corpus / 'epoch-{i}.conllu'.format(i=i)
        with nlp.use_params(optimizer.averages):
            parsed_docs, scores = evaluate(nlp, paths.dev.text, paths.dev.conllu, out_path)
            print_progress(i, losses, scores)
            _render_parses(i, parsed_docs[:50])

def _render_parses(i, to_render):
    to_render[0].user_data['title'] = "Batch %d" % i
    with Path('/tmp/parses.html').open('w') as file_:
        html = displacy.render(to_render[:5], style='dep', page=True)
        file_.write(html)

if __name__ == '__main__':
    plac.call(main)
<commit_msg>Fix vector loading for ud_train<commit_after>'''Train for CONLL 2017 UD treebank evaluation. Takes .conllu files, writes
.conllu format for development data, allowing the official scorer to be used.
'''
from __future__ import unicode_literals
import plac
import tqdm
from pathlib import Path
import re
import sys
import json

import spacy
import spacy.util
from ..tokens import Token, Doc
from ..gold import GoldParse
from ..util import compounding, minibatch_by_words
from ..syntax.nonproj import projectivize
from ..matcher import Matcher
from .. import displacy
from collections import defaultdict, Counter
from timeit import default_timer as timer

import itertools
import random
import numpy.random
import cytoolz

from . import conll17_ud_eval

from .. import lang
from .. import lang
from ..lang import zh
from ..lang import ja

lang.zh.Chinese.Defaults.use_jieba = False
lang.ja.Japanese.Defaults.use_janome = False

random.seed(0)
numpy.random.seed(0)

################
# Data reading #
################

space_re = re.compile('\s+')
def split_text(text):
    return [space_re.sub(' ', par.strip()) for par in text.split('\n\n')]

def read_data(nlp, conllu_file, text_file, raw_text=True, oracle_segments=False,
              max_doc_length=None, limit=None):
    '''Read the CONLLU format into (Doc, GoldParse) tuples. If raw_text=True,
    include Doc objects created using nlp.make_doc and then aligned against
    the gold-standard sequences. If oracle_segments=True, include Doc objects
    created from the gold-standard segments. At least one must be True.'''
    if not raw_text and not oracle_segments:
        raise ValueError("At least one of raw_text or oracle_segments must be True")
    paragraphs = split_text(text_file.read())
    conllu = read_conllu(conllu_file)
    # sd is spacy doc; cd is conllu doc
    # cs is conllu sent, ct is conllu token
    docs = []
    golds = []
    for doc_id, (text, cd) in enumerate(zip(paragraphs, conllu)):
        sent_annots = []
        for cs in cd:
            sent = defaultdict(list)
            for id_, word, lemma, pos, tag, morph, head, dep, _, space_after in cs:
                if '.' in id_:
                    continue
                if '-' in id_:
                    continue
                id_ = int(id_)-1
                head = int(head)-1 if head != '0' else id_
                sent['words'].append(word)
                sent['tags'].append(tag)
                sent['heads'].append(head)
                sent['deps'].append('ROOT' if dep == 'root' else dep)
                sent['spaces'].append(space_after == '_')
            sent['entities'] = ['-'] * len(sent['words'])
            sent['heads'], sent['deps'] = projectivize(sent['heads'],
                                                       sent['deps'])
            if oracle_segments:
                docs.append(Doc(nlp.vocab, words=sent['words'], spaces=sent['spaces']))
                golds.append(GoldParse(docs[-1], **sent))

            sent_annots.append(sent)
            if raw_text and max_doc_length and len(sent_annots) >= max_doc_length:
                doc, gold = _make_gold(nlp, None, sent_annots)
                sent_annots = []
                docs.append(doc)
                golds.append(gold)
                if limit and len(docs) >= limit:
                    return docs, golds

        if raw_text and sent_annots:
            doc, gold = _make_gold(nlp, None, sent_annots)
            docs.append(doc)
            golds.append(gold)
        if limit and len(docs) >= limit:
            return docs, golds
    return docs, golds

def read_conllu(file_):
    docs = []
    sent = []
    doc = []
    for line in file_:
        if line.startswith('# newdoc'):
            if doc:
                docs.append(doc)
            doc = []
        elif line.startswith('#'):
            continue
        elif not line.strip():
            if sent:
                doc.append(sent)
            sent = []
        else:
            sent.append(list(line.strip().split('\t')))
            if len(sent[-1]) != 10:
                print(repr(line))
                raise ValueError
    if sent:
        doc.append(sent)
    if doc:
        docs.append(doc)
    return docs

def _make_gold(nlp, text, sent_annots):
    # Flatten the conll annotations, and adjust the head indices
    flat = defaultdict(list)
    for sent in sent_annots:
        flat['heads'].extend(len(flat['words'])+head for head in sent['heads'])
        for field in ['words', 'tags', 'deps', 'entities', 'spaces']:
            flat[field].extend(sent[field])
    # Construct text if necessary
    assert len(flat['words']) == len(flat['spaces'])
    if text is None:
        text = ''.join(word+' '*space for word, space in zip(flat['words'], flat['spaces']))
    doc = nlp.make_doc(text)
    flat.pop('spaces')
    gold = GoldParse(doc, **flat)
    return doc, gold

#############################
# Data transforms for spaCy #
#############################

def golds_to_gold_tuples(docs, golds):
    '''Get out the annoying 'tuples' format used by begin_training, given the
    GoldParse objects.'''
    tuples = []
    for doc, gold in zip(docs, golds):
        text = doc.text
        ids, words, tags, heads, labels, iob = zip(*gold.orig_annot)
        sents = [((ids, words, tags, heads, labels, iob), [])]
        tuples.append((text, sents))
    return tuples

##############
# Evaluation #
##############

def evaluate(nlp, text_loc, gold_loc, sys_loc, limit=None):
    with text_loc.open('r', encoding='utf8') as text_file:
        texts = split_text(text_file.read())
        docs = list(nlp.pipe(texts))
    with sys_loc.open('w', encoding='utf8') as out_file:
        write_conllu(docs, out_file)
    with gold_loc.open('r', encoding='utf8') as gold_file:
        gold_ud = conll17_ud_eval.load_conllu(gold_file)
        with sys_loc.open('r', encoding='utf8') as sys_file:
            sys_ud = conll17_ud_eval.load_conllu(sys_file)
        scores = conll17_ud_eval.evaluate(gold_ud, sys_ud)
    return docs, scores

def write_conllu(docs, file_):
    merger = Matcher(docs[0].vocab)
    merger.add('SUBTOK', None, [{'DEP': 'subtok', 'op': '+'}])
    for i, doc in enumerate(docs):
        matches = merger(doc)
        spans = [doc[start:end+1] for _, start, end in matches]
        offsets = [(span.start_char, span.end_char) for span in spans]
        for start_char, end_char in offsets:
            doc.merge(start_char, end_char)
        file_.write("# newdoc id = {i}\n".format(i=i))
        for j, sent in enumerate(doc.sents):
            file_.write("# sent_id = {i}.{j}\n".format(i=i, j=j))
            file_.write("# text = {text}\n".format(text=sent.text))
            for k, token in enumerate(sent):
                file_.write(token._.get_conllu_lines(k) + '\n')
            file_.write('\n')

def print_progress(itn, losses, ud_scores):
    fields = {
        'dep_loss': losses.get('parser', 0.0),
        'tag_loss': losses.get('tagger', 0.0),
        'words': ud_scores['Words'].f1 * 100,
        'sents': ud_scores['Sentences'].f1 * 100,
        'tags': ud_scores['XPOS'].f1 * 100,
        'uas': ud_scores['UAS'].f1 * 100,
        'las': ud_scores['LAS'].f1 * 100,
    }
    header = ['Epoch', 'Loss', 'LAS', 'UAS', 'TAG', 'SENT', 'WORD']
    if itn == 0:
        print('\t'.join(header))
    tpl = '\t'.join((
        '{:d}',
        '{dep_loss:.1f}',
        '{las:.1f}',
        '{uas:.1f}',
        '{tags:.1f}',
        '{sents:.1f}',
        '{words:.1f}',
    ))
    print(tpl.format(itn, **fields))

#def get_sent_conllu(sent, sent_id):
#    lines = ["# sent_id = {sent_id}".format(sent_id=sent_id)]

def get_token_conllu(token, i):
    if token._.begins_fused:
        n = 1
        while token.nbor(n)._.inside_fused:
            n += 1
        id_ = '%d-%d' % (i, i+n)
        lines = [id_, token.text, '_', '_', '_', '_', '_', '_', '_', '_']
    else:
        lines = []
    if token.head.i == token.i:
        head = 0
    else:
        head = i + (token.head.i - token.i) + 1
    fields = [str(i+1), token.text, token.lemma_, token.pos_, token.tag_, '_',
              str(head), token.dep_.lower(), '_', '_']
    lines.append('\t'.join(fields))
    return '\n'.join(lines)

Token.set_extension('get_conllu_lines', method=get_token_conllu)
Token.set_extension('begins_fused', default=False)
Token.set_extension('inside_fused', default=False)

##################
# Initialization #
##################

def load_nlp(corpus, config):
    lang = corpus.split('_')[0]
    nlp = spacy.blank(lang)
    if config.vectors:
        nlp.vocab.from_disk(Path(config.vectors) / 'vocab')
    return nlp

def initialize_pipeline(nlp, docs, golds, config):
    nlp.add_pipe(nlp.create_pipe('parser'))
    if config.multitask_tag:
        nlp.parser.add_multitask_objective('tag')
    if config.multitask_sent:
        nlp.parser.add_multitask_objective('sent_start')
    nlp.add_pipe(nlp.create_pipe('tagger'))
    for gold in golds:
        for tag in gold.tags:
            if tag is not None:
                nlp.tagger.add_label(tag)
    return nlp.begin_training(lambda: golds_to_gold_tuples(docs, golds))

########################
# Command line helpers #
########################

class Config(object):
    def __init__(self, vectors=None, max_doc_length=10, multitask_tag=True,
            multitask_sent=True, nr_epoch=30, batch_size=1000, dropout=0.2):
        for key, value in locals().items():
            setattr(self, key, value)

    @classmethod
    def load(cls, loc):
        with Path(loc).open('r', encoding='utf8') as file_:
            cfg = json.load(file_)
        return cls(**cfg)

class Dataset(object):
    def __init__(self, path, section):
        self.path = path
        self.section = section
        self.conllu = None
        self.text = None
        for file_path in self.path.iterdir():
            name = file_path.parts[-1]
            if section in name and name.endswith('conllu'):
                self.conllu = file_path
            elif section in name and name.endswith('txt'):
                self.text = file_path
        if self.conllu is None:
            msg = "Could not find .txt file in {path} for {section}"
            raise IOError(msg.format(section=section, path=path))
        if self.text is None:
            msg = "Could not find .txt file in {path} for {section}"
        self.lang = self.conllu.parts[-1].split('-')[0].split('_')[0]

class TreebankPaths(object):
    def __init__(self, ud_path, treebank, **cfg):
        self.train = Dataset(ud_path / treebank, 'train')
        self.dev = Dataset(ud_path / treebank, 'dev')
        self.lang = self.train.lang

@plac.annotations(
    ud_dir=("Path to Universal Dependencies corpus", "positional", None, Path),
    corpus=("UD corpus to train and evaluate on, e.g. en, es_ancora, etc",
            "positional", None, str),
    parses_dir=("Directory to write the development parses", "positional", None, Path),
    config=("Path to json formatted config file", "positional"),
    limit=("Size limit", "option", "n", int)
)
def main(ud_dir, parses_dir, config, corpus, limit=0):
    config = Config.load(config)
    paths = TreebankPaths(ud_dir, corpus)
    if not (parses_dir / corpus).exists():
        (parses_dir / corpus).mkdir()
    print("Train and evaluate", corpus, "using lang", paths.lang)
    nlp = load_nlp(paths.lang, config)

    docs, golds = read_data(nlp, paths.train.conllu.open(), paths.train.text.open(),
                            max_doc_length=config.max_doc_length, limit=limit)

    optimizer = initialize_pipeline(nlp, docs, golds, config)

    batch_sizes = compounding(config.batch_size //10, config.batch_size, 1.001)
    for i in range(config.nr_epoch):
        docs = [nlp.make_doc(doc.text) for doc in docs]
        Xs = list(zip(docs, golds))
        random.shuffle(Xs)
        batches = minibatch_by_words(Xs, size=batch_sizes)
        losses = {}
        n_train_words = sum(len(doc) for doc in docs)
        with tqdm.tqdm(total=n_train_words, leave=False) as pbar:
            for batch in batches:
                batch_docs, batch_gold = zip(*batch)
                pbar.update(sum(len(doc) for doc in batch_docs))
                nlp.update(batch_docs, batch_gold, sgd=optimizer,
                           drop=config.dropout, losses=losses)

        out_path = parses_dir / corpus / 'epoch-{i}.conllu'.format(i=i)
        with nlp.use_params(optimizer.averages):
            parsed_docs, scores = evaluate(nlp, paths.dev.text, paths.dev.conllu, out_path)
            print_progress(i, losses, scores)
            _render_parses(i, parsed_docs[:50])

def _render_parses(i, to_render):
    to_render[0].user_data['title'] = "Batch %d" % i
    with Path('/tmp/parses.html').open('w') as file_:
        html = displacy.render(to_render[:5], style='dep', page=True)
        file_.write(html)

if __name__ == '__main__':
    plac.call(main)
<|endoftext|>


<commit_before>import collections
from baseline.embeddings import register_embeddings, create_embeddings
from eight_mile.pytorch.embeddings import *
from eight_mile.pytorch.serialize import load_tlm_npz, load_tlm_output_npz, tlm_load_state_dict
from eight_mile.utils import read_config_stream, mime_type
from baseline.vectorizers import load_bert_vocab

class PyTorchEmbeddingsModel(PyTorchEmbeddings):
    """A subclass of embeddings layers to prep them for registration and creation via baseline.

    In tensorflow this layer handles the creation of placeholders and things like that so the
    embeddings layer can just be tensor in tensor out but in pytorch all it does is strip the
    unused `name` input and register them.
    """
    def __init__(self, _=None, **kwargs):
        super().__init__(**kwargs)

    @classmethod
    def create(cls, model, name, **kwargs):
        kwargs.pop("dsz", None)
        return cls(name, vsz=model.vsz, dsz=model.dsz, weights=model.weights, **kwargs)

@register_embeddings(name='default')
class LookupTableEmbeddingsModel(PyTorchEmbeddingsModel, LookupTableEmbeddings):
    pass

@register_embeddings(name='char-conv')
class CharConvEmbeddingsModel(PyTorchEmbeddingsModel, CharConvEmbeddings):
    pass

@register_embeddings(name='char-lstm')
class CharLSTMEmbeddingsModel(PyTorchEmbeddingsModel, CharLSTMEmbeddings):
    pass

@register_embeddings(name='char-transformer')
class CharTransformerEmbeddingsModel(PyTorchEmbeddingsModel, CharTransformerEmbeddings):
    pass

@register_embeddings(name='positional')
class PositionalLookupTableEmbeddingsModel(PyTorchEmbeddingsModel, PositionalLookupTableEmbeddings):
    pass

@register_embeddings(name='learned-positional')
class LearnedPositionalLookupTableEmbeddingsModel(PyTorchEmbeddingsModel, LearnedPositionalLookupTableEmbeddings):
    pass

@register_embeddings(name='learned-positional-w-bias')
class LearnedPositionalLookupTableEmbeddingsWithBiasModel(PyTorchEmbeddingsModel, LearnedPositionalLookupTableEmbeddingsWithBias):
    pass

@register_embeddings(name='bert-lookup-table-embeddings')
class BERTLookupTableEmbeddingsModel(PyTorchEmbeddingsModel, BERTLookupTableEmbeddings):
    pass

@register_embeddings(name='positional-char-conv')
class PositionalCharConvEmbeddingsModel(PyTorchEmbeddingsModel, PositionalCharConvEmbeddings):
    pass

@register_embeddings(name='learned-positional-char-conv')
class LearnedPositionalCharConvEmbeddingsModel(PyTorchEmbeddingsModel, LearnedPositionalCharConvEmbeddings):
    pass

@register_embeddings(name='positional-char-lstm')
class PositionalCharLSTMEmbeddingsModel(PyTorchEmbeddingsModel, PositionalCharLSTMEmbeddings):
    pass

@register_embeddings(name='learned-positional-char-lstm')
class LearnedPositionalCharLSTMEmbeddingsModel(PyTorchEmbeddingsModel, LearnedPositionalCharLSTMEmbeddings):
    pass

class TransformerLMEmbeddings(PyTorchEmbeddings):
    """Support embeddings trained with the TransformerLanguageModel class

    This method supports either subword or word embeddings, not characters

    """
    def __init__(self, **kwargs):
        super().__init__()
        # You dont actually have to pass this if you are using the `load_bert_vocab` call from your
        # tokenizer.  In this case, a singleton variable will contain the vocab and it will be returned
        # by `load_bert_vocab`
        # If you trained your model with MEAD/Baseline, you will have a `*.json` file which would want to
        # reference here
        vocab_file = kwargs.get('vocab_file')
        if vocab_file and os.path.exists(vocab_file):
            if vocab_file.endswith('.json'):
                self.vocab = read_config_stream(kwargs.get('vocab_file'))
            else:
                self.vocab = load_bert_vocab(kwargs.get('vocab_file'))
        else:
            self.vocab = kwargs.get('vocab', kwargs.get('known_vocab'))
            if self.vocab is None or isinstance(self.vocab, collections.Counter):
                self.vocab = load_bert_vocab(None)
        # When we reload, allows skipping restoration of these embeddings
        # If the embedding wasnt trained with token types, this allows us to add them later
        self.skippable = set(listify(kwargs.get('skip_restore_embeddings', [])))

        self.cls_index = self.vocab.get('[CLS]', self.vocab.get('<s>'))
        self.vsz = max(self.vocab.values()) + 1
        self.d_model = int(kwargs.get('dsz', kwargs.get('d_model', 768)))
        self.init_embed(**kwargs)
        self.proj_to_dsz = pytorch_linear(self.dsz, self.d_model) if self.dsz != self.d_model else _identity
        self.init_transformer(**kwargs)

    @property
    def dsz(self):
        return self.embeddings.output_dim

    def embed(self, input, token_type):
        return self.embeddings({'x': input, 'tt': token_type})

    def init_embed(self, **kwargs):
        # If you are using BERT, you probably want to use either
        # `learned-positional` with a token type feature
        # or `learned-positional-w-bias` if you dont care about the token type
        embed_type = kwargs.get('word_embed_type', 'learned-positional')
        x_embedding = create_embeddings(vsz=self.vsz, dsz=self.d_model, embed_type=embed_type, offset=kwargs.get("offset", 0),
                                        mxlen=kwargs.get('mxlen', 512))

        embeddings = {'x': x_embedding}
        # This is for BERT support when we are using 2 features
        token_type_vsz = kwargs.get('token_type_vsz')
        if token_type_vsz:
            tt_unif = float(kwargs.get('token_type_embed_unif', 0.001))
            tt_embedding = LookupTableEmbeddings(weights=torch.randn((token_type_vsz, self.d_model))*tt_unif,
                                                 padding_idx=None)
            embeddings['tt'] = tt_embedding
        # For bert, make sure this is `sum-layer-norm`
        reduction = kwargs.get('embeddings_reduction', kwargs.get('reduction', 'sum'))
        embeddings_dropout = kwargs.get('embeddings_dropout', 0.1)
        self.embeddings = EmbeddingsStack(embeddings, dropout_rate=embeddings_dropout, reduction=reduction)

    def init_transformer(self, **kwargs):
        num_layers = int(kwargs.get('layers', 12))
        num_heads = int(kwargs.get('num_heads', 12))
        pdrop = kwargs.get('dropout', 0.1)
        ff_pdrop = kwargs.get('ffn_dropout', 0.1)
        d_ff = int(kwargs.get('d_ff', 3072))
        d_k = kwargs.get('d_k')
        rpr_k = kwargs.get('rpr_k')
        layer_norms_after = kwargs.get('layer_norms_after', False)
        layer_norm_eps = float(kwargs.get('layer_norm_eps', 1e-12))
        activation = kwargs.get('activation', 'gelu')
        windowed_ra = kwargs.get('windowed_ra', False)
        rpr_value_on = kwargs.get('rpr_value_on', True)
        is_mlp = kwargs.get("mlp", False)
        if is_mlp:
            self.transformer = GatedMLPEncoderStack(self.d_model, pdrop=pdrop, layers=num_layers,
                                                    nctx=kwargs.get('nctx', 256),
                                                    activation=activation, ffn_pdrop=ff_pdrop,
                                                    layer_norm_eps=layer_norm_eps)
        else:

            self.transformer = TransformerEncoderStack(num_heads, d_model=self.d_model, pdrop=pdrop, scale=True,
                                                       layers=num_layers, d_ff=d_ff, rpr_k=rpr_k, d_k=d_k,
                                                       activation=activation, ffn_pdrop=ff_pdrop,
                                                       layer_norms_after=layer_norms_after, layer_norm_eps=layer_norm_eps,
                                                       windowed_ra=windowed_ra, rpr_value_on=rpr_value_on)
        self.mlm = kwargs.get('mlm', True)
        self.finetune = kwargs.get('finetune', True)

    def forward(self, x, token_type=None):
        with torch.no_grad() if not self.finetune else contextlib.ExitStack():
            # the following line masks out the attention to padding tokens
            input_mask = torch.zeros(x.shape, device=x.device, dtype=torch.long).masked_fill(x != Offsets.PAD, 1).unsqueeze(1).unsqueeze(1)
            # A causal LM should have a subsequent mask; and a masked LM should have no mask
            if not self.mlm:
                input_mask = input_mask & subsequent_mask(x.shape[1]).type_as(input_mask)
            embedding = self.embed(x, token_type)
            embedding = self.proj_to_dsz(embedding)
            transformer_out = self.transformer((embedding, input_mask))
            z = self.get_output(x, transformer_out)
            return z

    def get_output(self, inputs, z):
        return z

    def get_vocab(self):
        return self.vocab

    def get_vsz(self):
        return self.vsz

    def get_dsz(self):
        return self.d_model

    @classmethod
    def load(cls, embeddings, **kwargs):
        c = cls("tlm-words-embed", **kwargs)

        if embeddings.endswith('.bin'):
            # HuggingFace checkpoint, convert on the fly
            from eight_mile.pytorch.serialize import load_tlm_transformers_bin, BERT_HF_FT_LAYER_MAP
            unmatch = load_tlm_transformers_bin(c, embeddings, replace_layers=BERT_HF_FT_LAYER_MAP)
            if unmatch['missing'] or unmatch['unexpected']:
                raise Exception("Unable to load the HuggingFace checkpoint")
        if mime_type(embeddings) == 'application/zip' and not embeddings.endswith("pth"):
            keys_to_restore = set(list(c.embeddings.keys()))
            filtered_keys = keys_to_restore.difference(c.skippable)
            if not keys_to_restore:
                raise Exception("No keys to restore!")
            if len(filtered_keys) < len(keys_to_restore):
                logger.warning("Restoring only key [%s]", ' '.join(filtered_keys))
            load_tlm_npz(c, embeddings, filtered_keys)
        else:
            tlm_load_state_dict(c, embeddings)
        return c

@register_embeddings(name='tlm-words-embed')
class TransformerLMEmbeddingsModel(PyTorchEmbeddingsModel, TransformerLMEmbeddings):
    """Register embedding model for usage in mead"""
    pass

def _identity(x):
    return x

def _mean_pool(inputs, embeddings):
    mask = (inputs != Offsets.PAD)
    seq_lengths = mask.sum(1).float()
    embeddings = embeddings.masked_fill(mask.unsqueeze(-1) == False, 0.)
    return embeddings.sum(1)/seq_lengths.unsqueeze(-1)

def _max_pool(inputs, embeddings):
    mask = (inputs != Offsets.PAD)
    embeddings = embeddings.masked_fill(mask.unsqueeze(-1) == False, 0.)
    return torch.max(embeddings, 1, False)[0]

@register_embeddings(name='tlm-words-embed-pooled')
class TransformerLMPooledEmbeddingsModel(TransformerLMEmbeddingsModel):

    def __init__(self, name, **kwargs):
        super().__init__(name=name, **kwargs)
        self.pooling = kwargs.get('pooling', 'cls')
        reduction_d_k = kwargs.get('reduction_d_k', self.d_model)
        dropout = kwargs.get('dropout', 0.1)  # use the same dropout as the transformer encoder for reduction layer
        if self.pooling == 'max':
            self.pooling_op = _max_pool
        elif self.pooling == 'mean':
            self.pooling_op = _mean_pool
        elif self.pooling == 'sqrt_length':
            self.pooling_op = self._sqrt_length_pool
        elif self.pooling == '2HA':
            self.reduction_layer = TwoHeadConcat(self.d_model, dropout, scale=False, d_k=reduction_d_k)
            self.pooling_op = self._att_reduction
        elif self.pooling == 'SHA':
            self.reduction_layer = SingleHeadReduction(self.d_model, dropout, scale=False, d_k=reduction_d_k)
            self.pooling_op = self._att_reduction
        else:
            self.pooling_op = self._cls_pool

    def _att_reduction(self, inputs, embeddings):
        mask = (inputs != Offsets.PAD)
        att_mask = mask.unsqueeze(1).unsqueeze(1)
        reduced = self.reduction_layer((embeddings, embeddings, embeddings, att_mask))
        return reduced

    def get_dsz(self):
        if self.pooling == '2HA':
            return 2*self.d_model
        return self.d_model

    def _sqrt_length_pool(self, inputs, embeddings):
        mask = (inputs != Offsets.PAD)
        lengths = mask.sum(1)
        sqrt_length = lengths.float()
        embeddings = embeddings.masked_fill(mask.unsqueeze(-1) == False, 0.)
        embeddings = embeddings.sum(1) * sqrt_length.sqrt().unsqueeze(-1)
        return embeddings

    def _cls_pool(self, inputs, tensor):
        # Would prefer
        # tensor[inputs == self.cls_index]
        # but ONNX export fails
        B = tensor.shape[0]
        mask = (inputs == self.cls_index).unsqueeze(-1).expand_as(tensor)
        pooled = tensor.masked_select(mask).view(B, -1)
        return pooled

    def get_output(self, inputs, z):
        z = self.pooling_op(inputs, z)
        return z

@register_embeddings(name='tlm-words-embed-pooled-output')
class TransformerLMPooledEmbeddingsWithOutputModel(TransformerLMPooledEmbeddingsModel):

    def __init__(self, name, **kwargs):
        super().__init__(name, **kwargs)
        self._output_dim = kwargs.get('output_dim', self.d_model)
        self.output_layer = pytorch_linear(self.d_model, self.output_dim)

    def get_dsz(self):
        return self._output_dim

    def get_output(self, inputs, z):
        z = self.pooling_op(inputs, z)
        return self.output_layer(z)

    @classmethod
    def load(cls, embeddings, **kwargs):
        c = cls("tlm-words-embed", **kwargs)

        if embeddings.endswith('.bin'):
            # HuggingFace checkpoint, convert on the fly
            from eight_mile.pytorch.serialize import load_tlm_transformers_bin, BERT_HF_FT_LAYER_MAP
            unmatch = load_tlm_transformers_bin(c, embeddings, replace_layers=BERT_HF_FT_LAYER_MAP)
            if unmatch['missing'] or unmatch['unexpected']:
                raise Exception("Unable to load the HuggingFace checkpoint")
        if mime_type(embeddings) == 'application/zip' and not embeddings.endswith("pth"):
            keys_to_restore = set(list(c.embeddings.keys()))
            filtered_keys = keys_to_restore.difference(c.skippable)
            if not keys_to_restore:
                raise Exception("No keys to restore!")
            if len(filtered_keys) < len(keys_to_restore):
                logger.warning("Restoring only key [%s]", ' '.join(filtered_keys))
            load_tlm_output_npz(c, embeddings, filtered_keys)
        else:
            tlm_load_state_dict(c, embeddings)
        return c

@register_embeddings(name='tlm-words-embed-pooled2d')
class TransformerLMPooled2DEmbeddingsModel(TransformerLMPooledEmbeddingsModel):

    def forward(self, xch, token_type=None):
        _0, _1, W = xch.shape
        pooled = super().forward(xch.view(-1, W))
        return pooled.view(_0, _1, self.get_dsz())
<commit_msg>add str_map to tlm-words-embed-pooled-output (#851)<commit_after>import collections
from baseline.embeddings import register_embeddings, create_embeddings
from eight_mile.pytorch.embeddings import *
from eight_mile.pytorch.serialize import load_tlm_npz, load_tlm_output_npz, tlm_load_state_dict
from eight_mile.utils import read_config_stream, mime_type
from baseline.vectorizers import load_bert_vocab

class PyTorchEmbeddingsModel(PyTorchEmbeddings):
    """A subclass of embeddings layers to prep them for registration and creation via baseline.

    In tensorflow this layer handles the creation of placeholders and things like that so the
    embeddings layer can just be tensor in tensor out but in pytorch all it does is strip the
    unused `name` input and register them.
    """
    def __init__(self, _=None, **kwargs):
        super().__init__(**kwargs)

    @classmethod
    def create(cls, model, name, **kwargs):
        kwargs.pop("dsz", None)
        return cls(name, vsz=model.vsz, dsz=model.dsz, weights=model.weights, **kwargs)

@register_embeddings(name='default')
class LookupTableEmbeddingsModel(PyTorchEmbeddingsModel, LookupTableEmbeddings):
    pass

@register_embeddings(name='char-conv')
class CharConvEmbeddingsModel(PyTorchEmbeddingsModel, CharConvEmbeddings):
    pass

@register_embeddings(name='char-lstm')
class CharLSTMEmbeddingsModel(PyTorchEmbeddingsModel, CharLSTMEmbeddings):
    pass

@register_embeddings(name='char-transformer')
class CharTransformerEmbeddingsModel(PyTorchEmbeddingsModel, CharTransformerEmbeddings):
    pass

@register_embeddings(name='positional')
class PositionalLookupTableEmbeddingsModel(PyTorchEmbeddingsModel, PositionalLookupTableEmbeddings):
    pass

@register_embeddings(name='learned-positional')
class LearnedPositionalLookupTableEmbeddingsModel(PyTorchEmbeddingsModel, LearnedPositionalLookupTableEmbeddings):
    pass

@register_embeddings(name='learned-positional-w-bias')
class LearnedPositionalLookupTableEmbeddingsWithBiasModel(PyTorchEmbeddingsModel, LearnedPositionalLookupTableEmbeddingsWithBias):
    pass

@register_embeddings(name='bert-lookup-table-embeddings')
class BERTLookupTableEmbeddingsModel(PyTorchEmbeddingsModel, BERTLookupTableEmbeddings):
    pass

@register_embeddings(name='positional-char-conv')
class PositionalCharConvEmbeddingsModel(PyTorchEmbeddingsModel, PositionalCharConvEmbeddings):
    pass

@register_embeddings(name='learned-positional-char-conv')
class LearnedPositionalCharConvEmbeddingsModel(PyTorchEmbeddingsModel, LearnedPositionalCharConvEmbeddings):
    pass

@register_embeddings(name='positional-char-lstm')
class PositionalCharLSTMEmbeddingsModel(PyTorchEmbeddingsModel, PositionalCharLSTMEmbeddings):
    pass

@register_embeddings(name='learned-positional-char-lstm')
class LearnedPositionalCharLSTMEmbeddingsModel(PyTorchEmbeddingsModel, LearnedPositionalCharLSTMEmbeddings):
    pass

class TransformerLMEmbeddings(PyTorchEmbeddings):
    """Support embeddings trained with the TransformerLanguageModel class

    This method supports either subword or word embeddings, not characters

    """
    def __init__(self, **kwargs):
        super().__init__()
        # You dont actually have to pass this if you are using the `load_bert_vocab` call from your
        # tokenizer.  In this case, a singleton variable will contain the vocab and it will be returned
        # by `load_bert_vocab`
        # If you trained your model with MEAD/Baseline, you will have a `*.json` file which would want to
        # reference here
        vocab_file = kwargs.get('vocab_file')
        if vocab_file and os.path.exists(vocab_file):
            if vocab_file.endswith('.json'):
                self.vocab = read_config_stream(kwargs.get('vocab_file'))
            else:
                self.vocab = load_bert_vocab(kwargs.get('vocab_file'))
        else:
            self.vocab = kwargs.get('vocab', kwargs.get('known_vocab'))
            if self.vocab is None or isinstance(self.vocab, collections.Counter):
                self.vocab = load_bert_vocab(None)
        # When we reload, allows skipping restoration of these embeddings
        # If the embedding wasnt trained with token types, this allows us to add them later
        self.skippable = set(listify(kwargs.get('skip_restore_embeddings', [])))

        self.cls_index = self.vocab.get('[CLS]', self.vocab.get('<s>'))
        self.vsz = max(self.vocab.values()) + 1
        self.d_model = int(kwargs.get('dsz', kwargs.get('d_model', 768)))
        self.init_embed(**kwargs)
        self.proj_to_dsz = pytorch_linear(self.dsz, self.d_model) if self.dsz != self.d_model else _identity
        self.init_transformer(**kwargs)

    @property
    def dsz(self):
        return self.embeddings.output_dim

    def embed(self, input, token_type):
        return self.embeddings({'x': input, 'tt': token_type})

    def init_embed(self, **kwargs):
        # If you are using BERT, you probably want to use either
        # `learned-positional` with a token type feature
        # or `learned-positional-w-bias` if you dont care about the token type
        embed_type = kwargs.get('word_embed_type', 'learned-positional')
        x_embedding = create_embeddings(vsz=self.vsz, dsz=self.d_model, embed_type=embed_type, offset=kwargs.get("offset", 0),
                                        mxlen=kwargs.get('mxlen', 512))

        embeddings = {'x': x_embedding}
        # This is for BERT support when we are using 2 features
        token_type_vsz = kwargs.get('token_type_vsz')
        if token_type_vsz:
            tt_unif = float(kwargs.get('token_type_embed_unif', 0.001))
            tt_embedding = LookupTableEmbeddings(weights=torch.randn((token_type_vsz, self.d_model))*tt_unif,
                                                 padding_idx=None)
            embeddings['tt'] = tt_embedding
        # For bert, make sure this is `sum-layer-norm`
        reduction = kwargs.get('embeddings_reduction', kwargs.get('reduction', 'sum'))
        embeddings_dropout = kwargs.get('embeddings_dropout', 0.1)
        self.embeddings = EmbeddingsStack(embeddings, dropout_rate=embeddings_dropout, reduction=reduction)

    def init_transformer(self, **kwargs):
        num_layers = int(kwargs.get('layers', 12))
        num_heads = int(kwargs.get('num_heads', 12))
        pdrop = kwargs.get('dropout', 0.1)
        ff_pdrop = kwargs.get('ffn_dropout', 0.1)
        d_ff = int(kwargs.get('d_ff', 3072))
        d_k = kwargs.get('d_k')
        rpr_k = kwargs.get('rpr_k')
        layer_norms_after = kwargs.get('layer_norms_after', False)
        layer_norm_eps = float(kwargs.get('layer_norm_eps', 1e-12))
        activation = kwargs.get('activation', 'gelu')
        windowed_ra = kwargs.get('windowed_ra', False)
        rpr_value_on = kwargs.get('rpr_value_on', True)
        is_mlp = kwargs.get("mlp", False)
        if is_mlp:
            self.transformer = GatedMLPEncoderStack(self.d_model, pdrop=pdrop, layers=num_layers,
                                                    nctx=kwargs.get('nctx', 256),
                                                    activation=activation, ffn_pdrop=ff_pdrop,
                                                    layer_norm_eps=layer_norm_eps)
        else:

            self.transformer = TransformerEncoderStack(num_heads, d_model=self.d_model, pdrop=pdrop, scale=True,
                                                       layers=num_layers, d_ff=d_ff, rpr_k=rpr_k, d_k=d_k,
                                                       activation=activation, ffn_pdrop=ff_pdrop,
                                                       layer_norms_after=layer_norms_after, layer_norm_eps=layer_norm_eps,
                                                       windowed_ra=windowed_ra, rpr_value_on=rpr_value_on)
        self.mlm = kwargs.get('mlm', True)
        self.finetune = kwargs.get('finetune', True)

    def forward(self, x, token_type=None):
        with torch.no_grad() if not self.finetune else contextlib.ExitStack():
            # the following line masks out the attention to padding tokens
            input_mask = torch.zeros(x.shape, device=x.device, dtype=torch.long).masked_fill(x != Offsets.PAD, 1).unsqueeze(1).unsqueeze(1)
            # A causal LM should have a subsequent mask; and a masked LM should have no mask
            if not self.mlm:
                input_mask = input_mask & subsequent_mask(x.shape[1]).type_as(input_mask)
            embedding = self.embed(x, token_type)
            embedding = self.proj_to_dsz(embedding)
            transformer_out = self.transformer((embedding, input_mask))
            z = self.get_output(x, transformer_out)
            return z

    def get_output(self, inputs, z):
        return z

    def get_vocab(self):
        return self.vocab

    def get_vsz(self):
        return self.vsz

    def get_dsz(self):
        return self.d_model

    @classmethod
    def load(cls, embeddings, **kwargs):
        c = cls("tlm-words-embed", **kwargs)

        if embeddings.endswith('.bin'):
            # HuggingFace checkpoint, convert on the fly
            from eight_mile.pytorch.serialize import load_tlm_transformers_bin, BERT_HF_FT_LAYER_MAP
            unmatch = load_tlm_transformers_bin(c, embeddings, replace_layers=BERT_HF_FT_LAYER_MAP)
            if unmatch['missing'] or unmatch['unexpected']:
                raise Exception("Unable to load the HuggingFace checkpoint")
        if mime_type(embeddings) == 'application/zip' and not embeddings.endswith("pth"):
            keys_to_restore = set(list(c.embeddings.keys()))
            filtered_keys = keys_to_restore.difference(c.skippable)
            if not keys_to_restore:
                raise Exception("No keys to restore!")
            if len(filtered_keys) < len(keys_to_restore):
                logger.warning("Restoring only key [%s]", ' '.join(filtered_keys))
            load_tlm_npz(c, embeddings, filtered_keys)
        else:
            tlm_load_state_dict(c, embeddings)
        return c

@register_embeddings(name='tlm-words-embed')
class TransformerLMEmbeddingsModel(PyTorchEmbeddingsModel, TransformerLMEmbeddings):
    """Register embedding model for usage in mead"""
    pass

def _identity(x):
    return x

def _mean_pool(inputs, embeddings):
    mask = (inputs != Offsets.PAD)
    seq_lengths = mask.sum(1).float()
    embeddings = embeddings.masked_fill(mask.unsqueeze(-1) == False, 0.)
    return embeddings.sum(1)/seq_lengths.unsqueeze(-1)

def _max_pool(inputs, embeddings):
    mask = (inputs != Offsets.PAD)
    embeddings = embeddings.masked_fill(mask.unsqueeze(-1) == False, 0.)
    return torch.max(embeddings, 1, False)[0]

@register_embeddings(name='tlm-words-embed-pooled')
class TransformerLMPooledEmbeddingsModel(TransformerLMEmbeddingsModel):

    def __init__(self, name, **kwargs):
        super().__init__(name=name, **kwargs)
        self.pooling = kwargs.get('pooling', 'cls')
        reduction_d_k = kwargs.get('reduction_d_k', self.d_model)
        dropout = kwargs.get('dropout', 0.1)  # use the same dropout as the transformer encoder for reduction layer
        if self.pooling == 'max':
            self.pooling_op = _max_pool
        elif self.pooling == 'mean':
            self.pooling_op = _mean_pool
        elif self.pooling == 'sqrt_length':
            self.pooling_op = self._sqrt_length_pool
        elif self.pooling == '2HA':
            self.reduction_layer = TwoHeadConcat(self.d_model, dropout, scale=False, d_k=reduction_d_k)
            self.pooling_op = self._att_reduction
        elif self.pooling == 'SHA':
            self.reduction_layer = SingleHeadReduction(self.d_model, dropout, scale=False, d_k=reduction_d_k)
            self.pooling_op = self._att_reduction
        else:
            self.pooling_op = self._cls_pool

    def _att_reduction(self, inputs, embeddings):
        mask = (inputs != Offsets.PAD)
        att_mask = mask.unsqueeze(1).unsqueeze(1)
        reduced = self.reduction_layer((embeddings, embeddings, embeddings, att_mask))
        return reduced

    def get_dsz(self):
        if self.pooling == '2HA':
            return 2*self.d_model
        return self.d_model

    def _sqrt_length_pool(self, inputs, embeddings):
        mask = (inputs != Offsets.PAD)
        lengths = mask.sum(1)
        sqrt_length = lengths.float()
        embeddings = embeddings.masked_fill(mask.unsqueeze(-1) == False, 0.)
        embeddings = embeddings.sum(1) * sqrt_length.sqrt().unsqueeze(-1)
        return embeddings

    def _cls_pool(self, inputs, tensor):
        # Would prefer
        # tensor[inputs == self.cls_index]
        # but ONNX export fails
        B = tensor.shape[0]
        mask = (inputs == self.cls_index).unsqueeze(-1).expand_as(tensor)
        pooled = tensor.masked_select(mask).view(B, -1)
        return pooled

    def get_output(self, inputs, z):
        z = self.pooling_op(inputs, z)
        return z

@register_embeddings(name='tlm-words-embed-pooled-output')
class TransformerLMPooledEmbeddingsWithOutputModel(TransformerLMPooledEmbeddingsModel):

    def __init__(self, name, **kwargs):
        super().__init__(name, **kwargs)
        self._output_dim = kwargs.get('output_dim', self.d_model)
        self.output_layer = pytorch_linear(self.d_model, self.output_dim)

    def get_dsz(self):
        return self._output_dim

    def get_output(self, inputs, z):
        z = self.pooling_op(inputs, z)
        return self.output_layer(z)

    @classmethod
    def load(cls, embeddings, **kwargs):
        c = cls("tlm-words-embed", **kwargs)

        if embeddings.endswith('.bin'):
            # HuggingFace checkpoint, convert on the fly
            from eight_mile.pytorch.serialize import load_tlm_transformers_bin, BERT_HF_FT_LAYER_MAP
            unmatch = load_tlm_transformers_bin(c, embeddings, replace_layers=BERT_HF_FT_LAYER_MAP)
            if unmatch['missing'] or unmatch['unexpected']:
                raise Exception("Unable to load the HuggingFace checkpoint")
        if mime_type(embeddings) == 'application/zip' and not embeddings.endswith("pth"):
            keys_to_restore = set(list(c.embeddings.keys()))
            filtered_keys = keys_to_restore.difference(c.skippable)
            if not keys_to_restore:
                raise Exception("No keys to restore!")
            if len(filtered_keys) < len(keys_to_restore):
                logger.warning("Restoring only key [%s]", ' '.join(filtered_keys))
            load_tlm_output_npz(c, embeddings, filtered_keys)
        else:
            tlm_load_state_dict(c, embeddings,
                                str_map={'model.embeddings.embeddings.0.':'', 'model.output_layer': 'output_layer'})
        return c

@register_embeddings(name='tlm-words-embed-pooled2d')
class TransformerLMPooled2DEmbeddingsModel(TransformerLMPooledEmbeddingsModel):

    def forward(self, xch, token_type=None):
        _0, _1, W = xch.shape
        pooled = super().forward(xch.view(-1, W))
        return pooled.view(_0, _1, self.get_dsz())
<|endoftext|>


<commit_before># Copyright 2018 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Writes a Perf-formated json file with stats about the given web file."""

import json
import os
import subprocess
import sys

def main():
  input_file = sys.argv[1]
  out_dir = sys.argv[2]
  keystr = sys.argv[3]
  propstr = sys.argv[4]

  results = {
    'keys': { },
    'results': { }
  }

  props = propstr.split(' ')
  for i in range(0, len(props), 2):
    results[props[i]] = props[i+1]

  keys = keystr.split(' ')
  for i in range(0, len(keys), 2):
    results['keys'][keys[i]] = keys[i+1]

  r = {
    'total_size_bytes': os.path.getsize(input_file)
  }

  # Make a copy to avoid destroying the hardlinked file.
  # Swarming hardlinks in the builds from isolated cache.
  temp_file = input_file + '_tmp'
  subprocess.check_call(['cp', input_file, temp_file])
  subprocess.check_call(['gzip', temp_file])

  r['gzip_size_bytes'] = os.path.getsize(temp_file + '.gz')

  name = os.path.basename(input_file)

  results['results'][name] = {
    # We need this top level layer 'config'/slice
    # Other analysis methods (e.g. libskia) might have
    # slices for data on the 'code' section, etc.
    'default' : r,
  }

  # Make debugging easier
  print json.dumps(results, indent=2)

  with open(os.path.join(out_dir, name+'.json'), 'w') as output:
    output.write(json.dumps(results, indent=2))

if '__main__' == __name__:
  main()
<commit_msg>Fix key in perf data<commit_after># Copyright 2018 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Writes a Perf-formated json file with stats about the given web file."""

import json
import os
import subprocess
import sys

def main():
  input_file = sys.argv[1]
  out_dir = sys.argv[2]
  keystr = sys.argv[3]
  propstr = sys.argv[4]

  results = {
    'key': { },
    'results': { }
  }

  props = propstr.split(' ')
  for i in range(0, len(props), 2):
    results[props[i]] = props[i+1]

  keys = keystr.split(' ')
  for i in range(0, len(keys), 2):
    results['key'][keys[i]] = keys[i+1]

  r = {
    'total_size_bytes': os.path.getsize(input_file)
  }

  # Make a copy to avoid destroying the hardlinked file.
  # Swarming hardlinks in the builds from isolated cache.
  temp_file = input_file + '_tmp'
  subprocess.check_call(['cp', input_file, temp_file])
  subprocess.check_call(['gzip', temp_file])

  r['gzip_size_bytes'] = os.path.getsize(temp_file + '.gz')

  name = os.path.basename(input_file)

  results['results'][name] = {
    # We need this top level layer 'config'/slice
    # Other analysis methods (e.g. libskia) might have
    # slices for data on the 'code' section, etc.
    'default' : r,
  }

  # Make debugging easier
  print json.dumps(results, indent=2)

  with open(os.path.join(out_dir, name+'.json'), 'w') as output:
    output.write(json.dumps(results, indent=2))

if '__main__' == __name__:
  main()
<|endoftext|>


<commit_before>"""
desisim.scripts.brightsims
==========================

Generate a canonical set of bright-time simulations.
"""

from __future__ import division, print_function

import os
import sys
import numpy as np
import argparse
from argparse import Namespace

from astropy.io import fits
from astropy.table import Table

import desisim.scripts.quickbrick as quickbrick
from desispec.log import get_logger, DEBUG
from desispec.io.util import write_bintable, makepath

def parse(options=None):
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                     description='Generate a canonical set of bright-time simulations.')

    parser.add_argument('-b', '--brickname', type=str, help='unique output brickname suffix',
                        default='brightsims', metavar='')

    parser.add_argument('--objtype', type=str,  help='BGS, MWS, or BRIGHT_MIX', default='BRIGHT_MIX', metavar='')
    parser.add_argument('--nbrick', type=int,  help='number of bricks to simulate', default=10, metavar='')
    parser.add_argument('--nspec', type=int,  help='number of spectra (per brick) to simulate', default=100, metavar='')

    parser.add_argument('-s', '--seed', type=int,  help='random seed', default=None, metavar='')
    parser.add_argument('-o', '--outdir', type=str,  help='output directory for simulation parameters',
                        default='.', metavar='')
    parser.add_argument('--brickdir', type=str,  help='top-level output bricks directory',
                        default='./bricks', metavar='')
    parser.add_argument('-v', '--verbose', action='store_true', help='toggle on verbose output')

    parser.add_argument('--exptime-range', type=float, default=(300,300), nargs=2, metavar='',
                        help='minimum and maximum exposure time (s)')
    parser.add_argument('--airmass-range', type=float, default=(1.25,1.25), nargs=2, metavar='',
                        help='minimum and maximum airmass')
    parser.add_argument('--moon-phase-range', type=float, default=(0.0,1.0), nargs=2, metavar='',
                        help='minimum and maximum lunar phase (0=full, 1=new')
    parser.add_argument('--moon-angle-range', type=float, default=(0,90), nargs=2, metavar='',
                        help='minimum and maximum lunar separation angle (0-90 deg')
    parser.add_argument('--moon-zenith-range', type=float, default=(0,60), nargs=2, metavar='',
                        help='minimum and maximum lunar zenith angle (0-90 deg')

    args = None
    if options is None:
        args = parser.parse_args()
    else:
        args = parser.parse_args(options)
    return args

def main(args):

    # Set up the logger.
    if args.verbose:
        log = get_logger(DEBUG)
    else:
        log = get_logger()

    objtype = args.objtype.upper()

    log.debug('Using OBJTYPE {}'.format(objtype))
    log.debug('Simulating {:g} bricks each with {:g} spectra'.format(args.nbrick, args.nspec))

    # Draw priors uniformly given the input ranges.
    rand = np.random.RandomState(args.seed)
    exptime = rand.uniform(args.exptime_range[0], args.exptime_range[1], args.nbrick)
    airmass = rand.uniform(args.airmass_range[0], args.airmass_range[1], args.nbrick)
    moonphase = rand.uniform(args.moon_phase_range[0], args.moon_phase_range[1], args.nbrick)
    moonangle = rand.uniform(args.moon_angle_range[0], args.moon_angle_range[1], args.nbrick)
    moonzenith = rand.uniform(args.moon_zenith_range[0], args.moon_zenith_range[1], args.nbrick)

    # Build a metadata table with the simulation inputs.
    metafile = makepath(os.path.join(args.outdir, '{}-input.fits'.format(args.brickname)))
    metacols = [
        ('BRICKNAME', 'S20'),
        ('SEED', 'S20'),
        ('EXPTIME', 'f4'),
        ('AIRMASS', 'f4'),
        ('MOONPHASE', 'f4'),
        ('MOONANGLE', 'f4'),
        ('MOONZENITH', 'f4')]
    meta = Table(np.zeros(args.nbrick, dtype=metacols))
    meta['EXPTIME'].unit = 's'
    meta['MOONANGLE'].unit = 'deg'
    meta['MOONZENITH'].unit = 'deg'

    meta['BRICKNAME'] = ['{}-{:03d}'.format(args.brickname, ii) for ii in range(args.nbrick)]
    meta['EXPTIME'] = exptime
    meta['AIRMASS'] = airmass
    meta['MOONPHASE'] = moonphase
    meta['MOONANGLE'] = moonangle
    meta['MOONZENITH'] = moonzenith

    log.debug('Writing {}'.format(metafile))
    write_bintable(metafile, meta, extname='METADATA', clobber=True)

    for ii in range(args.nbrick):
        thisbrick = meta['BRICKNAME'][ii]
        log.debug('Building brick {}'.format(thisbrick))

        brickargs = ['--brickname', thisbrick,
                     '--objtype', args.objtype,
                     '--nspec', '{}'.format(args.nspec),
                     '--outdir', args.brickdir,
                     '--outdir-truth', args.brickdir,
                     '--exptime', '{}'.format(exptime[ii]),
                     '--airmass', '{}'.format(airmass[ii]),
                     '--moon-phase', '{}'.format(moonphase[ii]),
                     '--moon-angle', '{}'.format(moonangle[ii]),
                     '--moon-zenith', '{}'.format(moonzenith[ii])]
        if args.seed is not None:
            brickargs.append('--seed')
            brickargs.append('{}'.format(args.seed))

        quickargs = quickbrick.parse(brickargs)
        if args.verbose:
            quickargs.verbose = True
        quickbrick.main(quickargs)
<commit_msg>a bit more organization by bricks<commit_after>"""
desisim.scripts.brightsims
==========================

Generate a canonical set of bright-time simulations.
"""

from __future__ import division, print_function

import os
import sys
import numpy as np
import argparse
from argparse import Namespace

from astropy.io import fits
from astropy.table import Table

import desisim.scripts.quickbrick as quickbrick
from desispec.log import get_logger, DEBUG
from desispec.io.util import write_bintable, makepath

def parse(options=None):
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                     description='Generate a canonical set of bright-time simulations.')

    parser.add_argument('-b', '--brickname', type=str, help='unique output brickname suffix',
                        default='brightsims', metavar='')

    parser.add_argument('--objtype', type=str,  help='BGS, MWS, or BRIGHT_MIX', default='BRIGHT_MIX', metavar='')
    parser.add_argument('--nbrick', type=int,  help='number of bricks to simulate', default=10, metavar='')
    parser.add_argument('--nspec', type=int,  help='number of spectra (per brick) to simulate', default=100, metavar='')

    parser.add_argument('-s', '--seed', type=int,  help='random seed', default=None, metavar='')
    parser.add_argument('-o', '--outdir', type=str,  help='output directory for simulation parameters',
                        default='.', metavar='')
    parser.add_argument('--brickdir', type=str,  help='top-level output bricks directory',
                        default='./bricks', metavar='')
    parser.add_argument('-v', '--verbose', action='store_true', help='toggle on verbose output')

    parser.add_argument('--exptime-range', type=float, default=(300,300), nargs=2, metavar='',
                        help='minimum and maximum exposure time (s)')
    parser.add_argument('--airmass-range', type=float, default=(1.25,1.25), nargs=2, metavar='',
                        help='minimum and maximum airmass')
    parser.add_argument('--moon-phase-range', type=float, default=(0.0,1.0), nargs=2, metavar='',
                        help='minimum and maximum lunar phase (0=full, 1=new')
    parser.add_argument('--moon-angle-range', type=float, default=(0,90), nargs=2, metavar='',
                        help='minimum and maximum lunar separation angle (0-90 deg')
    parser.add_argument('--moon-zenith-range', type=float, default=(0,60), nargs=2, metavar='',
                        help='minimum and maximum lunar zenith angle (0-90 deg')

    args = None
    if options is None:
        args = parser.parse_args()
    else:
        args = parser.parse_args(options)
    return args

def main(args):

    # Set up the logger.
    if args.verbose:
        log = get_logger(DEBUG)
    else:
        log = get_logger()

    objtype = args.objtype.upper()

    log.debug('Using OBJTYPE {}'.format(objtype))
    log.debug('Simulating {:g} bricks each with {:g} spectra'.format(args.nbrick, args.nspec))

    # Draw priors uniformly given the input ranges.
    rand = np.random.RandomState(args.seed)
    exptime = rand.uniform(args.exptime_range[0], args.exptime_range[1], args.nbrick)
    airmass = rand.uniform(args.airmass_range[0], args.airmass_range[1], args.nbrick)
    moonphase = rand.uniform(args.moon_phase_range[0], args.moon_phase_range[1], args.nbrick)
    moonangle = rand.uniform(args.moon_angle_range[0], args.moon_angle_range[1], args.nbrick)
    moonzenith = rand.uniform(args.moon_zenith_range[0], args.moon_zenith_range[1], args.nbrick)

    # Build a metadata table with the simulation inputs.
    metafile = makepath(os.path.join(args.outdir, '{}-input.fits'.format(args.brickname)))
    metacols = [
        ('BRICKNAME', 'S20'),
        ('SEED', 'S20'),
        ('EXPTIME', 'f4'),
        ('AIRMASS', 'f4'),
        ('MOONPHASE', 'f4'),
        ('MOONANGLE', 'f4'),
        ('MOONZENITH', 'f4')]
    meta = Table(np.zeros(args.nbrick, dtype=metacols))
    meta['EXPTIME'].unit = 's'
    meta['MOONANGLE'].unit = 'deg'
    meta['MOONZENITH'].unit = 'deg'

    meta['BRICKNAME'] = ['{}-{:03d}'.format(args.brickname, ii) for ii in range(args.nbrick)]
    meta['EXPTIME'] = exptime
    meta['AIRMASS'] = airmass
    meta['MOONPHASE'] = moonphase
    meta['MOONANGLE'] = moonangle
    meta['MOONZENITH'] = moonzenith

    log.debug('Writing {}'.format(metafile))
    write_bintable(metafile, meta, extname='METADATA', clobber=True)

    for ii in range(args.nbrick):
        thisbrick = meta['BRICKNAME'][ii]
        log.debug('Building brick {}'.format(thisbrick))

        brickargs = ['--brickname', thisbrick,
                     '--objtype', args.objtype,
                     '--nspec', '{}'.format(args.nspec),
                     '--outdir', os.path.join(args.brickdir, thisbrick),
                     '--outdir-truth', os.path.join(args.brickdir, thisbrick),
                     '--exptime', '{}'.format(exptime[ii]),
                     '--airmass', '{}'.format(airmass[ii]),
                     '--moon-phase', '{}'.format(moonphase[ii]),
                     '--moon-angle', '{}'.format(moonangle[ii]),
                     '--moon-zenith', '{}'.format(moonzenith[ii])]
        if args.seed is not None:
            brickargs.append('--seed')
            brickargs.append('{}'.format(args.seed))

        quickargs = quickbrick.parse(brickargs)
        if args.verbose:
            quickargs.verbose = True
        quickbrick.main(quickargs)
<|endoftext|>


<commit_before>from __future__ import division, print_function
from keylevel import _KeyLevel
from helpers import ReaderError, SUPPRESS, Namespace

__all__ = ['InputReader', 'ReaderError', 'SUPPRESS']

class InputReader(_KeyLevel):
    '''\
    :py:class:`InputReader` is a class that is designed to read in
    an input file and return the information contained based on rules
    supplied to the class using a syntax similar to what is used in
    the :py:mod:`argparse` module in the Python standard library.

    :py:class:`InputReader` accepts blocks-type, line-type and
    boolean-type keys, mutually exclusive groups, required keys,
    defaults, and more.

    :keyword comment:
        Defines what is a comment in the input block.
        This can be a single string or a list of strings.
        The default is *['#']*.
    :type comment: str or list, optional
    :keyword case:
        Tells if the keys are case-sensitive or not.
        The default is :py:const:`False`.
    :type case: bool, optional
    :keyword ignoreunknown:
        Ignore keys that are not defined.  The default is
        :py:const:`False`.
    :type ignoreunknown: bool, optional
    :keyword default:
        The default default that will be given when a
        key is created without a default.
    :type default: optional
    '''

    def __init__(self, comment=['#'], case=False, ignoreunknown=False,
                 default=None):
        '''Initiallize the :py:class:`InputReader` class.'''
        _KeyLevel.__init__(self, case=case)

        # Name (main is default)
        self.name = 'main'

        # What constitutes a comment?
        if isinstance(comment, str):
            comment = [comment]
        self._comment = comment
        try:
            for x in self._comment:
                if not isinstance(x, str):
                    raise ValueError ('comment value must be a str, '
                                      'given '+repr(x))
        except TypeError:
            raise ValueError ('comment value must be a str, '
                              'given '+repr(self._comment))

        # Ignore unknown keys?
        self._ignoreunknown = ignoreunknown
        if not isinstance(self._ignoreunknown, bool):
            raise ValueError ('ignoreunknown value must be a bool, '
                              'given '+repr(self._ignoreunknown))

        # The default default
        self._default = default

    def read_input(self, filename):
        '''\
        Reads in the input from a given file using the supplied rules.

        :argument filename:
            The name of the file to read in, or a StringIO or list of strings
            containing the input itself.
        :type filename: str
        :rtype: :py:class:`Namespace`: This class contains the read-in data
            each key is stored as members of the class.
        :exception:
            :py:exc:`ReaderError`: Any known errors will be raised with
            this custom exception.
        '''

        # Read in the file, removeing comments and extra whitespace/newlines
        f = self._read_in_file(filename)

        # Parse this key level, recursively reading lower levels
        i, namespace = self._parse_key_level(f, 0)
        namespace.__filename__ = filename
        namespace._order.append('__filename__')
        namespace.__input__ = f
        namespace._order.append('__input__')
        return namespace

    def _read_in_file(self, filename):
        '''Store the filename as a list'''

        f = []
        # Assume a filename was given
        try:
            fl = [x.rstrip() for x in open(filename)]
        except (IOError, OSError) as e:
            raise ReaderError ('Cannot read in file '+filename+':'+str(e))
        except TypeError:
            # Assume a StringIO object was given
            try:
                fl = filename.getvalue().split('\n')
            except AttributeError:
                # Assume an iterable of strings was given
                try:
                    fl = [x.rstrip() for x in filename]
                except AttributeError:
                    raise ValueError ('Unknown object passed to '
                                      'read_input: '+repr(filename))

        # Read in the data
        for line in fl:
            # Remove comments
            for com in self._comment:
                if com in line:
                    line = line.partition(com)[0]
            # Add this to the list
            f.append(line.strip())

        return f

<commit_msg>Removed original input file from Namespace<commit_after>from __future__ import division, print_function
from keylevel import _KeyLevel
from helpers import ReaderError, SUPPRESS, Namespace

__all__ = ['InputReader', 'ReaderError', 'SUPPRESS']

class InputReader(_KeyLevel):
    '''\
    :py:class:`InputReader` is a class that is designed to read in
    an input file and return the information contained based on rules
    supplied to the class using a syntax similar to what is used in
    the :py:mod:`argparse` module in the Python standard library.

    :py:class:`InputReader` accepts blocks-type, line-type and
    boolean-type keys, mutually exclusive groups, required keys,
    defaults, and more.

    :keyword comment:
        Defines what is a comment in the input block.
        This can be a single string or a list of strings.
        The default is *['#']*.
    :type comment: str or list, optional
    :keyword case:
        Tells if the keys are case-sensitive or not.
        The default is :py:const:`False`.
    :type case: bool, optional
    :keyword ignoreunknown:
        Ignore keys that are not defined.  The default is
        :py:const:`False`.
    :type ignoreunknown: bool, optional
    :keyword default:
        The default default that will be given when a
        key is created without a default.
    :type default: optional
    '''

    def __init__(self, comment=['#'], case=False, ignoreunknown=False,
                 default=None):
        '''Initiallize the :py:class:`InputReader` class.'''
        _KeyLevel.__init__(self, case=case)

        # Name (main is default)
        self.name = 'main'

        # What constitutes a comment?
        if isinstance(comment, str):
            comment = [comment]
        self._comment = comment
        try:
            for x in self._comment:
                if not isinstance(x, str):
                    raise ValueError ('comment value must be a str, '
                                      'given '+repr(x))
        except TypeError:
            raise ValueError ('comment value must be a str, '
                              'given '+repr(self._comment))

        # Ignore unknown keys?
        self._ignoreunknown = ignoreunknown
        if not isinstance(self._ignoreunknown, bool):
            raise ValueError ('ignoreunknown value must be a bool, '
                              'given '+repr(self._ignoreunknown))

        # The default default
        self._default = default

    def read_input(self, filename):
        '''\
        Reads in the input from a given file using the supplied rules.

        :argument filename:
            The name of the file to read in, or a StringIO or list of strings
            containing the input itself.
        :type filename: str
        :rtype: :py:class:`Namespace`: This class contains the read-in data
            each key is stored as members of the class.
        :exception:
            :py:exc:`ReaderError`: Any known errors will be raised with
            this custom exception.
        '''

        # Read in the file, removeing comments and extra whitespace/newlines
        f = self._read_in_file(filename)

        # Parse this key level, recursively reading lower levels
        i, namespace = self._parse_key_level(f, 0)
        return namespace

    def _read_in_file(self, filename):
        '''Store the filename as a list'''

        f = []
        # Assume a filename was given
        try:
            fl = [x.rstrip() for x in open(filename)]
        except (IOError, OSError) as e:
            raise ReaderError ('Cannot read in file '+filename+':'+str(e))
        except TypeError:
            # Assume a StringIO object was given
            try:
                fl = filename.getvalue().split('\n')
            except AttributeError:
                # Assume an iterable of strings was given
                try:
                    fl = [x.rstrip() for x in filename]
                except AttributeError:
                    raise ValueError ('Unknown object passed to '
                                      'read_input: '+repr(filename))

        # Read in the data
        for line in fl:
            # Remove comments
            for com in self._comment:
                if com in line:
                    line = line.partition(com)[0]
            # Add this to the list
            f.append(line.strip())

        return f

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Track your time with Timewarrior.

Timewarrior is a time tracking utility that offers simple stopwatch features
as well as sophisticated calendar-base backfill, along with flexible reporting.
See https://taskwarrior.org/docs/timewarrior for more information.

Configuration parameters:
    cache_timeout: refresh interval for this module, otherwise auto
        (default None)
    filter: specify interval and/or tag to filter (default None)
    format: display format for this module
        (default '[Timew {format_time}]|No Timew')
    format_datetime: specify strftime characters to format (default {})
    format_duration: display format for time duration
        (default '\?not_zero [{days}d ][{hours}:]{minutes}:{seconds}')
    format_tag: display format for tags (default '\?color=state_tag {name}')
    format_tag_separator: show separator if more than one (default ' ')
    format_time: display format for tracked times
        (default '[\?color=state_time [{format_tag} ]{format_duration}]')
    format_time_separator: show separator if more than one (default ' ')
    thresholds: specify color thresholds to use
        *(default {'state_tag': [(0, 'darkgray'), (1, 'darkgray')],
        'state_time': [(0, 'darkgray'), (1, 'degraded')]})*

Format placeholders:
    {format_time} format for tracked times
    {tracking} time tracking state, eg False, True

format_time placeholders:
    {state} time tracking state, eg False, True
    {format_tag} format for tags
    {format_duration} format for time duration
    {start} start date, eg 20171021T010203Z
    {end} end date, eg 20171021T010203Z

format_tag placeholders:
    {name} tag name, eg gaming, studying, gardening

format_datetime placeholders:
    key: start, end
    value: strftime characters, eg '%b %d' ----> 'Oct 06'

format_duration placeholders:
    {days} days
    {hours} hours
    {minutes} minutes
    {seconds} seconds

Color thresholds:
    format_time:
        state_time: print color based on the state of time tracking
    format_tag:
        state_tag:  print color based on the state of time tracking

Requires:
    timew: feature-rich time tracking utility

Recommendations:
    We can refresh a module using `py3-cmd` command.
    An excellent example of using this command in a function.

    ```
    ~/.{bash,zsh}{rc,_profile}
    ---------------------------
    function timew () {
        command timew "$@" && py3-cmd refresh timewarrior
    }
    ```

    With this, you can consider giving `cache_timeout` a much larger number,
    eg 3600 (an hour), so the module does not need to be updated that often.

Examples:
```
# show times matching the filter, see documentation for more filters
timewarrior {
    filter = ':day'           # filter times not in 24 hours of current day
    filter = '12hours'        # filter times not in 12 hours of current time
    filter = '5min'           # filter times not in 5 minutes of current time
    filter = '1sec'           # filter times not in 1 second of current time
    filter = '5pm to 11:59pm  # filter times not in 5pm to 11:59pm range
}

# intervals
timewarrior {
    # if you are printing other intervals too with '1day' filter or so,
    # then you may want to add this too for better bar readability
    format_time_separator = ', '

    # you also can change the thresholds with different colors
    thresholds = {
        'state_tag': [(0, 'darkgray'), (1, 'degraded')],
        'state_time': [(0, 'darkgray'), (1, 'degraded')],
    }
}

# cache_timeout
timewarrior {
    # auto refresh every 10 seconds when there is no active time tracking
    # auto refresh every second when there is active time tracking
    cache_timeout = None

    # refresh every minute when there is no active time tracking
    # refresh every second when there is active time tracking
    cache_timeout = 60

    # explicit refresh every 20 seconds when there is no active time tracking
    # explicit refresh every 5 seconds when there is active time tracking
    cache_timeout = (20, 5)
}

# add your snippets here
timewarrior {
    format = "..."
}
```

@author lasers

SAMPLE OUTPUT
[
    {'full_text': 'Timew '},
    {'full_text': 'gaming ', 'color': '#a9a9a9'},
    {'full_text': '15:02 ', 'color': '#a9a9a9'},
    {'full_text': 'studying ', 'color': '#a9a9a9'},
    {'full_text': '03:42', 'color': '#ffff00'}
]

no_tag
[
    {'full_text': 'Timew '},
    {'full_text': 'gardening ', 'color': '#a9a9a9'},
    {'full_text': '20:37', 'color': '#ffff00'}
]

no_timew
{'full_text': 'No Timew'}
"""

from pytz import utc
from json import loads as json_loads
from datetime import datetime
from dateutil.relativedelta import relativedelta

STRING_NOT_INSTALLED = 'not installed'
DATETIME = '%Y%m%dT%H%M%SZ'
STRING_INVALID_TIMEOUT = 'invalid cache_timeout'

class Py3status:
    """
    """
    cache_timeout = None
    filter = None
    format = '[Timew {format_time}]|No Timew'
    format_datetime = {}
    format_duration = '\?not_zero [{days}d ][{hours}:]{minutes}:{seconds}'
    format_tag = '\?color=state_tag {name}'
    format_tag_separator = ' '
    format_time = '[\?color=state_time [{format_tag} ]{format_duration}]'
    format_time_separator = ' '
    thresholds = {
        'state_tag': [(0, 'darkgray'), (1, 'darkgray')],
        'state_time': [(0, 'darkgray'), (1, 'degraded')],
    }

    class Meta:
        update_config = {
            'update_placeholder_format': [
                {
                    'placeholder_formats': {
                        'minutes': ':02d',
                        'seconds': ':02d',
                    },
                    'format_strings': ['format_duration'],
                }
            ],
        }

    def post_config_hook(self):
        if not self.py3.check_commands('timew'):
            raise Exception(STRING_NOT_INSTALLED)

        if self.cache_timeout is None:
            self.sleep_timeout = 10
            self.cache_timeout = 0
        elif isinstance(self.cache_timeout, tuple):
            if len(self.cache_timeout) != 2:
                raise Exception(STRING_INVALID_TIMEOUT)
            self.sleep_timeout = self.cache_timeout[0]
            self.cache_timeout = self.cache_timeout[1]
        elif isinstance(self.cache_timeout, int):
            self.sleep_timeout = self.cache_timeout
            self.cache_timeout = 0

        self.timewarrior_command = 'timew export'
        if self.filter:
            self.timewarrior_command += ' {}'.format(self.filter)

        self.init = {'datetimes': []}
        for word in ['start', 'end']:
            if (self.py3.format_contains(self.format_time, word)) and (
                    word in self.format_datetime):
                self.init['datetimes'].append(word)

        self.tracking = None
        self.thresholds_init = {}
        for name in ('format', 'format_tag', 'format_time'):
            self.thresholds_init[name] = self.py3.get_color_names_list(
                getattr(self, name, '')
            )

    def _get_timewarrior_data(self):
        return json_loads(self.py3.command_output(self.timewarrior_command))

    def _manipulate(self, data):
        new_time = []
        self.tracking = False

        for time_index, time in zip(range(len(data), -1, -1), data):
            time['index'] = time_index
            time['state_time'] = 'end' not in time

            # tags
            new_tag = []
            time['tags'] = time.get('tags', [])
            for tag_name in time['tags']:
                tag_data = {'name': tag_name, 'state_tag': time['state_time']}
                for x in self.thresholds_init['format_tag']:
                    if x in tag_data:
                        self.py3.threshold_get_color(tag_data[x], x)
                new_tag.append(self.py3.safe_format(self.format_tag, tag_data))

            format_tag_separator = self.py3.safe_format(self.format_tag_separator)
            format_tag = self.py3.composite_join(format_tag_separator, new_tag)

            time['format_tag'] = format_tag
            del time['tags']

            # duraton
            if time['state_time']:
                self.tracking = True
                end = datetime.now(utc).utcnow()
            else:
                end = datetime.strptime(time['end'], DATETIME)

            start = datetime.strptime(time['start'], DATETIME)
            duration = relativedelta(end, start)

            time['format_duration'] = self.py3.safe_format(
                self.format_duration, {
                    'days': duration.days,
                    'hours': duration.hours,
                    'minutes': duration.minutes,
                    'seconds': duration.seconds,
                }
            )

            # datetime
            for word in self.init['datetimes']:
                if word in time:
                    time[word] = self.py3.safe_format(datetime.strftime(
                        datetime.strptime(time[word], DATETIME),
                        self.format_datetime[word]))

            # time
            for x in self.thresholds_init['format_time']:
                if x in time:
                    self.py3.threshold_get_color(time[x], x)

            new_time.append(self.py3.safe_format(self.format_time, time))

        format_time_separator = self.py3.safe_format(self.format_time_separator)
        format_time = self.py3.composite_join(format_time_separator, new_time)
        return format_time

    def timewarrior(self):
        timewarrior_data = self._get_timewarrior_data()
        format_time = self._manipulate(timewarrior_data)

        if self.tracking:
            cached_until = self.cache_timeout
        else:
            cached_until = self.sleep_timeout

        timew_data = {
            'format_time': format_time,
            'tracking': self.tracking,
        }

        for x in self.thresholds_init['format']:
            if x in timew_data:
                self.py3.threshold_get_color(timew_data[x], x)

        return {
            'cached_until': self.py3.time_in(cached_until),
            'full_text': self.py3.safe_format(self.format, timew_data)
        }

if __name__ == '__main__':
    from py3status.module_test import module_test
    module_test(Py3status)
<commit_msg>timewarrior: add '1day' filter to be less irritating<commit_after># -*- coding: utf-8 -*-
"""
Track your time with Timewarrior.

Timewarrior is a time tracking utility that offers simple stopwatch features
as well as sophisticated calendar-base backfill, along with flexible reporting.
See https://taskwarrior.org/docs/timewarrior for more information.

Configuration parameters:
    cache_timeout: refresh interval for this module, otherwise auto
        (default None)
    filter: specify interval and/or tag to filter (default '1day')
    format: display format for this module
        (default '[Timew {format_time}]|No Timew')
    format_datetime: specify strftime characters to format (default {})
    format_duration: display format for time duration
        (default '\?not_zero [{days}d ][{hours}:]{minutes}:{seconds}')
    format_tag: display format for tags (default '\?color=state_tag {name}')
    format_tag_separator: show separator if more than one (default ' ')
    format_time: display format for tracked times
        (default '[\?color=state_time [{format_tag} ]{format_duration}]')
    format_time_separator: show separator if more than one (default ' ')
    thresholds: specify color thresholds to use
        *(default {'state_tag': [(0, 'darkgray'), (1, 'darkgray')],
        'state_time': [(0, 'darkgray'), (1, 'degraded')]})*

Format placeholders:
    {format_time} format for tracked times
    {tracking} time tracking state, eg False, True

format_time placeholders:
    {state} time tracking state, eg False, True
    {format_tag} format for tags
    {format_duration} format for time duration
    {start} start date, eg 20171021T010203Z
    {end} end date, eg 20171021T010203Z

format_tag placeholders:
    {name} tag name, eg gaming, studying, gardening

format_datetime placeholders:
    key: start, end
    value: strftime characters, eg '%b %d' ----> 'Oct 06'

format_duration placeholders:
    {days} days
    {hours} hours
    {minutes} minutes
    {seconds} seconds

Color thresholds:
    format_time:
        state_time: print color based on the state of time tracking
    format_tag:
        state_tag:  print color based on the state of time tracking

Requires:
    timew: feature-rich time tracking utility

Recommendations:
    We can refresh a module using `py3-cmd` command.
    An excellent example of using this command in a function.

    ```
    ~/.{bash,zsh}{rc,_profile}
    ---------------------------
    function timew () {
        command timew "$@" && py3-cmd refresh timewarrior
    }
    ```

    With this, you can consider giving `cache_timeout` a much larger number,
    eg 3600 (an hour), so the module does not need to be updated that often.

Examples:
```
# show times matching the filter, see documentation for more filters
timewarrior {
    filter = ':day'           # filter times not in 24 hours of current day
    filter = '12hours'        # filter times not in 12 hours of current time
    filter = '5min'           # filter times not in 5 minutes of current time
    filter = '1sec'           # filter times not in 1 second of current time
    filter = '5pm to 11:59pm  # filter times not in 5pm to 11:59pm range
}

# intervals
timewarrior {
    # if you are printing other intervals too with '1day' filter or so,
    # then you may want to add this too for better bar readability
    format_time_separator = ', '

    # you also can change the thresholds with different colors
    thresholds = {
        'state_tag': [(0, 'darkgray'), (1, 'degraded')],
        'state_time': [(0, 'darkgray'), (1, 'degraded')],
    }
}

# cache_timeout
timewarrior {
    # auto refresh every 10 seconds when there is no active time tracking
    # auto refresh every second when there is active time tracking
    cache_timeout = None

    # refresh every minute when there is no active time tracking
    # refresh every second when there is active time tracking
    cache_timeout = 60

    # explicit refresh every 20 seconds when there is no active time tracking
    # explicit refresh every 5 seconds when there is active time tracking
    cache_timeout = (20, 5)
}

# add your snippets here
timewarrior {
    format = "..."
}
```

@author lasers

SAMPLE OUTPUT
[
    {'full_text': 'Timew '},
    {'full_text': 'gaming ', 'color': '#a9a9a9'},
    {'full_text': '15:02 ', 'color': '#a9a9a9'},
    {'full_text': 'studying ', 'color': '#a9a9a9'},
    {'full_text': '03:42', 'color': '#ffff00'}
]

no_tag
[
    {'full_text': 'Timew '},
    {'full_text': 'gardening ', 'color': '#a9a9a9'},
    {'full_text': '20:37', 'color': '#ffff00'}
]

no_timew
{'full_text': 'No Timew'}
"""

from pytz import utc
from json import loads as json_loads
from datetime import datetime
from dateutil.relativedelta import relativedelta

STRING_NOT_INSTALLED = 'not installed'
DATETIME = '%Y%m%dT%H%M%SZ'
STRING_INVALID_TIMEOUT = 'invalid cache_timeout'

class Py3status:
    """
    """
    cache_timeout = None
    filter = '1day'
    format = '[Timew {format_time}]|No Timew'
    format_datetime = {}
    format_duration = '\?not_zero [{days}d ][{hours}:]{minutes}:{seconds}'
    format_tag = '\?color=state_tag {name}'
    format_tag_separator = ' '
    format_time = '[\?color=state_time [{format_tag} ]{format_duration}]'
    format_time_separator = ' '
    thresholds = {
        'state_tag': [(0, 'darkgray'), (1, 'darkgray')],
        'state_time': [(0, 'darkgray'), (1, 'degraded')],
    }

    class Meta:
        update_config = {
            'update_placeholder_format': [
                {
                    'placeholder_formats': {
                        'minutes': ':02d',
                        'seconds': ':02d',
                    },
                    'format_strings': ['format_duration'],
                }
            ],
        }

    def post_config_hook(self):
        if not self.py3.check_commands('timew'):
            raise Exception(STRING_NOT_INSTALLED)

        if self.cache_timeout is None:
            self.sleep_timeout = 10
            self.cache_timeout = 0
        elif isinstance(self.cache_timeout, tuple):
            if len(self.cache_timeout) != 2:
                raise Exception(STRING_INVALID_TIMEOUT)
            self.sleep_timeout = self.cache_timeout[0]
            self.cache_timeout = self.cache_timeout[1]
        elif isinstance(self.cache_timeout, int):
            self.sleep_timeout = self.cache_timeout
            self.cache_timeout = 0

        self.timewarrior_command = 'timew export'
        if self.filter:
            self.timewarrior_command += ' {}'.format(self.filter)

        self.init = {'datetimes': []}
        for word in ['start', 'end']:
            if (self.py3.format_contains(self.format_time, word)) and (
                    word in self.format_datetime):
                self.init['datetimes'].append(word)

        self.tracking = None
        self.thresholds_init = {}
        for name in ('format', 'format_tag', 'format_time'):
            self.thresholds_init[name] = self.py3.get_color_names_list(
                getattr(self, name, '')
            )

    def _get_timewarrior_data(self):
        return json_loads(self.py3.command_output(self.timewarrior_command))

    def _manipulate(self, data):
        new_time = []
        self.tracking = False

        for time_index, time in zip(range(len(data), -1, -1), data):
            time['index'] = time_index
            time['state_time'] = 'end' not in time

            # tags
            new_tag = []
            time['tags'] = time.get('tags', [])
            for tag_name in time['tags']:
                tag_data = {'name': tag_name, 'state_tag': time['state_time']}
                for x in self.thresholds_init['format_tag']:
                    if x in tag_data:
                        self.py3.threshold_get_color(tag_data[x], x)
                new_tag.append(self.py3.safe_format(self.format_tag, tag_data))

            format_tag_separator = self.py3.safe_format(self.format_tag_separator)
            format_tag = self.py3.composite_join(format_tag_separator, new_tag)

            time['format_tag'] = format_tag
            del time['tags']

            # duraton
            if time['state_time']:
                self.tracking = True
                end = datetime.now(utc).utcnow()
            else:
                end = datetime.strptime(time['end'], DATETIME)

            start = datetime.strptime(time['start'], DATETIME)
            duration = relativedelta(end, start)

            time['format_duration'] = self.py3.safe_format(
                self.format_duration, {
                    'days': duration.days,
                    'hours': duration.hours,
                    'minutes': duration.minutes,
                    'seconds': duration.seconds,
                }
            )

            # datetime
            for word in self.init['datetimes']:
                if word in time:
                    time[word] = self.py3.safe_format(datetime.strftime(
                        datetime.strptime(time[word], DATETIME),
                        self.format_datetime[word]))

            # time
            for x in self.thresholds_init['format_time']:
                if x in time:
                    self.py3.threshold_get_color(time[x], x)

            new_time.append(self.py3.safe_format(self.format_time, time))

        format_time_separator = self.py3.safe_format(self.format_time_separator)
        format_time = self.py3.composite_join(format_time_separator, new_time)
        return format_time

    def timewarrior(self):
        timewarrior_data = self._get_timewarrior_data()
        format_time = self._manipulate(timewarrior_data)

        if self.tracking:
            cached_until = self.cache_timeout
        else:
            cached_until = self.sleep_timeout

        timew_data = {
            'format_time': format_time,
            'tracking': self.tracking,
        }

        for x in self.thresholds_init['format']:
            if x in timew_data:
                self.py3.threshold_get_color(timew_data[x], x)

        return {
            'cached_until': self.py3.time_in(cached_until),
            'full_text': self.py3.safe_format(self.format, timew_data)
        }

if __name__ == '__main__':
    from py3status.module_test import module_test
    module_test(Py3status)
<|endoftext|>