#!/usr/bin/env python3

# The main LemmaTag code can be found below.

# For an up-to-date version of this code, see https://github.com/Hyperparticle/LemmaTag

# Author: Daniel Kondratyuk, Tomáš Gavenčiak, and Milan Straka

import numpy as np

import tensorflow as tf

import argparse

import datetime

import os

import shutil

import sys

from tqdm import tqdm

from tensorflow.python.client import timeline

import logging

from logging import warning, info, debug, error

from util import morpho_dataset

from util.utils import MorphoAnalyzer, Tee, log_time, find_first, AddInputsWrapper

from util.tags import WholeTags, CharTags, DictTags

from model.encoder import encoder_network

from model.tag_decoder import tag_decoder, tag_features

from model.lemma_decoder import lemma_decoder, sense_predictor

class LemmaTagNetwork:

return lemmas, tags

if __name__ == "__main__":

# Parse arguments

parser = argparse.ArgumentParser()

# General and training arguments

parser.add_argument("--batch_size", default=32, type=int, help="Batch size.")

parser.add_argument("--epochs", default=40, type=int, help="Number of epochs.")

parser.add_argument("--threads", default=4, type=int, help="Maximum number of threads to use.")

parser.add_argument("--name", default="", type=str, help="Any name comment.")

parser.add_argument("--checkpoint", default="", type=str, help="Checkpoint restore directory.")

parser.add_argument("--beta_2", default=0.99, type=float, help="Adam beta 2.")

parser.add_argument("--learning_rate", default=0.001, type=float, help="Learning rate.")

parser.add_argument("--drop_rate_after", default=20, type=int, help="Number of epochs after which the rate is quartered every 10 epochs.")

parser.add_argument("--grad_clip", default=3.0, type=float, help="Gradient clipping (if set).")

parser.add_argument("--record_trace", default=False, action="store_true", help="Record training trace as Chrome trace (load at 'chrome://tracing/').")

parser.add_argument("--no_save_net", default=False, action="store_true", help="Skip checkoint saving (to save space when debugging).")

parser.add_argument("--only_eval", default=False, action="store_true", help="Skip training and only evaluate once (from a checkpoint).")

# Data and seed

parser.add_argument("--seed", default=42, type=int, help="Random seed.")

parser.add_argument("--train", default="data/sample-cs-cltt-ud-train.txt", type=str, help="Training data path.")

parser.add_argument("--dev", default="data/sample-cs-cltt-ud-dev.txt", type=str, help="Validation data path.")

parser.add_argument("--test", default="data/sample-cs-cltt-ud-test.txt", type=str, help="Test data path.")

parser.add_argument("--conllu", default=False, action="store_true", help="Using a conllu-formatted dataset")

parser.add_argument("--analyzer", default=None, type=str, help="Analyzer text file (default none).")

parser.add_argument("--max_sentences", default=None, type=int, help="Max sentences to load (for quick testing).")

# Dimensions and features

parser.add_argument("--cle_dim", default=64, type=int, help="Character-level embedding dimension.")

parser.add_argument("--rnn_cell", default="LSTM", type=str, help="RNN cell type.")

parser.add_argument("--rnn_cell_dim", default=128, type=int, help="RNN cell dimension.")

parser.add_argument("--rnn_layers", default=2, type=int, help="RNN layers.")

parser.add_argument("--we_dim", default=128, type=int, help="Word embedding dimension.")

parser.add_argument("--att_dim", default=64, type=int, help="Attention dimension.")

parser.add_argument("--predict_sense", default=False, action="store_true", help="Train and predict the sense as a part of the lemma (use dataset with separated sense for that).")

parser.add_argument("--no_tags_to_lemmas", default=False, action="store_true", help="Don't use tag components as a signal to the lemmatizer.")

parser.add_argument("--separate_rnn", default=False, action="store_true", help="Use separate RNN for tags and lemmas/senses.")

parser.add_argument("--separate_embed", default=False, action="store_true", help="Use separate embeddings for tags and lemmas/senses. Implies separate_rnn.")

parser.add_argument("--beams", default=None, type=int, help="Use beam search with the given no of beams.")

parser.add_argument("--loss_sense_w", default=0.1, type=float, help="Sense loss weight (if sense is separate).")

parser.add_argument("--loss_lem_w", default=1.0, type=float, help="Lemmatization loss weight.")

# Regularization

parser.add_argument("--dropout", default=0.5, type=float, help="Dropout rate")

parser.add_argument("--label_smoothing", default=0.1, type=float, help="Label smoothing.")

parser.add_argument("--lem_smoothing", default=0.0, type=float, help="Lemma label smoothing.")

parser.add_argument("--sense_smoothing", default=0.05, type=float, help="Sense label smoothing.")

parser.add_argument("--word_dropout", default=0.25, type=float, help="Word dropout")

parser.add_argument("--tag_signal_dropout", default=None, type=float, help="Tag signal dropout to lemmatizer")

parser.add_argument("--tag_type", default="char", choices=["char", "dict", "whole"], help="Compositional tag type.")

parser.add_argument("--compositional_tags_regularization", default=0.1, type=float, help="Compositional tags regularization.")

parser.add_argument("--whole_tags_regularization", default=1.0, type=float, help="Whole tags regularization.")

parser.add_argument("--beam_len_penalty", default=0.2, type=float, help="BeamSearch length_penalty_weight param.")

args = parser.parse_args()

if args.separate_embed:

args.separate_rnn = True

if args.only_eval:

args.epochs = 1

# Fix the random seed

np.random.seed(args.seed)

# Logdir, copy the source and log the outputs

if not os.path.exists("logs"): os.mkdir("logs")

basename = "LT-{}-{}-S{}".format(

datetime.datetime.now().strftime("%Y%m%d_%H%M%S"),

args.name, args.seed, )

args.logdir = "logs/" + basename

os.mkdir(args.logdir)

shutil.copy(__file__, args.logdir + "/taglem.py")

tee = Tee(args.logdir + "/log.txt")

tee.start()

args.realstderr = tee.stderr # A bit hacky ... (to leave progress bars out of logs)

logging.basicConfig(format='%(asctime)s [%(levelname)s] %(message)s', level=logging.DEBUG)

info("Running in {} with args: {}".format(args.logdir, str(args)))

info("Commandline: {}".format(' '.join(sys.argv)))

# Load the data

with log_time("load inputs"):

args.max_dev_sentences = args.max_sentences // 5 if args.max_sentences else None

train = morpho_dataset.MorphoDataset(args.train, max_sentences=args.max_sentences, conllu_format=args.conllu)

dev = morpho_dataset.MorphoDataset(args.dev, train=train, shuffle_batches=False, max_sentences=args.max_dev_sentences, conllu_format=args.conllu)

test = morpho_dataset.MorphoDataset(args.test, train=train, shuffle_batches=False, max_sentences=args.max_dev_sentences, conllu_format=args.conllu)

# analyzer = MorphoAnalyzer(args.analyzer) if args.analyzer else None

# Construct the network

if args.tag_type == "char":

args.tags = CharTags(train, args.compositional_tags_regularization, args.whole_tags_regularization)

elif args.tag_type == "dict":

# TODO: add support for dictionary tags

raise ValueError("Tag type not supported: " + args.tag_type)

# args.tags = DictTags(train, args.compositional_tags_regularization, args.whole_tags_regularization)

elif args.tag_type == "whole":

args.tags = WholeTags(train)

else:

raise ValueError("Invalid tag_type")

network = LemmaTagNetwork(threads=args.threads, seed=args.seed)

network.construct(args, len(train.factors[train.FORMS].words), len(train.factors[train.FORMS].alphabet),

len(train.factors[train.LEMMAS].alphabet), args.tags.num_tags(),

len(train.factors[train.SENSES].words), train.factors[train.LEMMAS].alphabet_map["<bow>"],

train.factors[train.LEMMAS].alphabet_map["<eow>"])

if args.checkpoint:

network.saver.restore(network.session, args.checkpoint)

# Train

dev_best = 0

for ep in range(args.epochs):

rate = args.learning_rate

if args.drop_rate_after and args.drop_rate_after <= ep:

rate = args.learning_rate * 0.25 ** (1 + ((ep - args.drop_rate_after) // 10))

if not args.only_eval:

info("Training epoch %d with rate %f", ep, rate)

network.train_epoch(train, args, rate=rate)

info("Evaluating dev")

dev_acc_tag, dev_acc_lem, dev_acc_lemsense = network.evaluate("dev", dev, args)

info(".. epoch {} (step {}) dev accuracy: {:.2f} tag, {:.2f} lemma, {:.2f} lemma with sense".format(

ep, network.session.run(network.global_step), 100 * dev_acc_tag, 100 * dev_acc_lem, 100 * dev_acc_lemsense))

if dev_acc_tag + dev_acc_lemsense > dev_best or ep == args.epochs - 1:

if not args.no_save_net and not args.only_eval: # To speed up testing / save disk space :)

network.saver.save(network.session, "{}/checkpoint".format(args.logdir), global_step=network.global_step, write_meta_graph=False)

for dset, name in [(dev, "dev"), (test, "test")]:

fname = "{}/taglem_{}_ep{}.txt".format(args.logdir, name, ep)

info("Predicting %s into %s", name, fname)

with open(fname, "w", encoding="utf-8") as ofile:

forms = dset.factors[dset.FORMS].strings

lemmas, tags = network.predict(dset, args)

for s in range(len(forms)):

for i in range(len(forms[s])):

print("{}\t{}\t{}".format(forms[s][i], lemmas[s][i], tags[s][i]), file=ofile)

print("", file=ofile)

dev_best = max(dev_best, dev_acc_tag + dev_acc_lemsense)