seqtagger/train_seq_tagger.py

import os
from flair.data import Corpus
from flair.datasets import ColumnCorpus
from flair.embeddings import *
from flair.trainers import ModelTrainer
from typing import List
from flair.models import SequenceTagger
from flair.trainers import ModelTrainer
import ast
from torch.utils.tensorboard import SummaryWriter
import yaml
import json
from shutil import copyfile, move, rmtree, copy2
from extract_score import extract_score

pram_file = "./params/seqtag_params.yaml"
metrics_file = "./metrics/seqtag_metrics.json"
plots_file = "./metrics"
tensorboard_folder = "tensorboard/"


# import params
with open(pram_file, 'r') as fd:
    params = yaml.safe_load(fd)


# 1. get the corpus

# define columns
# columns = {0: 'token', 1: 'tag', 2: 'space-after'}
columns = params["corpus"]["columns"]

# this is the folder in which train, test and dev files reside
data_folder = params["corpus"]["data_folder"]

train_file = params["corpus"]["train_file"]
test_file = params["corpus"]["test_file"]
dev_file = params["corpus"]["dev_file"]

# init a corpus using column format, data folder and the names of the train, dev and test files
corpus: Corpus = ColumnCorpus(data_folder, columns,
                              train_file=train_file,
                              test_file=test_file,
                              dev_file=dev_file)

# 2. what tag do we want to predict?
tag_type = 'ner'

# 3. make the tag dictionary from the corpus
tag_dictionary = corpus.make_tag_dictionary(tag_type=tag_type)

# 4. initialize embeddings
embedding_types: List[TokenEmbeddings] =[]

if "flair-backward" in params["embeddings"].keys():
    embedding_types.append(FlairEmbeddings(params["embeddings"]["flair-backward"], chars_per_chunk=512))

if "fasttext" in params["embeddings"].keys():
    embedding_types.append(WordEmbeddings(params["embeddings"]["fasttext"]))

embeddings: StackedEmbeddings = StackedEmbeddings(embeddings=embedding_types)

# 5. initialize sequence tagger
tagger: SequenceTagger = SequenceTagger(embeddings=embeddings,
                                        tag_dictionary=tag_dictionary,
                                        tag_type=tag_type,
                                        **params["seqtagger"]
                                        )

# 6. initialize trainer
trainer: ModelTrainer = ModelTrainer(tagger, corpus, **params["trainer"])

path = params["train"]["base_path"]

if params["type"]["is_fresh_start"]:
    try:
        rmtree(path, ignore_errors=True)
    except OSError as e:
        print("Error: %s : %s" % (path, e.strerror))

    try:
        rmtree(tensorboard_folder, ignore_errors=True)
    except OSError as e:
        print("Error: %s : %s" % (tensorboard_folder+"*", e.strerror))
# 7. start training
trainer.train(**params["train"])

# 8. eval
# run evaluation procedure

result, _ = tagger.evaluate(corpus.test, out_path=f"{path}/predictions.txt")
print(result.log_line)
jsonfile = open(metrics_file, 'w')
det_res = extract_score(result)
json.dump(det_res, jsonfile)

copy2(path+"loss.tsv", plots_file)
move("./runs/", tensorboard_folder + params["type"]["run_name"]+"/")