flair_embeddings_cli.py

import os.path
from flair.data import Dictionary
from flair.models import LanguageModel
from flair.trainers.language_model_trainer import LanguageModelTrainer, TextCorpus


def train_flair_embeddings(
    corpus_path="/data/ubertext/for_flair",
    dictionary_path="/home/dima/Projects/flair_embeddings/flair_dictionary.pkl",
    lm_file="./language_model_forward_no_amp_accum_grad_fixed",
    is_forward_lm=True,
    hidden_size=1024,
    sequence_length=250,
    mini_batch_size=256,
    max_epochs=20,
    gpus=1,
):
    # load the default character dictionary
    dictionary: Dictionary = Dictionary.load(dictionary_path)

    # get your corpus, process forward and at the character level
    corpus = TextCorpus(corpus_path, dictionary, is_forward_lm, character_level=True)

    checkpoint = os.path.join(lm_file, "checkpoint.pt")

    if os.path.exists(checkpoint):
        trainer = LanguageModelTrainer.load_checkpoint(checkpoint, corpus)
    else:
        # train your language model
        language_model = LanguageModel(dictionary, is_forward_lm, hidden_size=hidden_size, nlayers=1)
        trainer = LanguageModelTrainer(language_model, corpus)

    trainer.train(
        lm_file,
        sequence_length=sequence_length,
        mini_batch_size=mini_batch_size,
        max_epochs=max_epochs,
        use_amp=False,
        checkpoint=True,
    )


train_flair_embeddings()

flair_embeddings_cli_backwards.py

import os.path
from flair.data import Dictionary
from flair.models import LanguageModel
from flair.trainers.language_model_trainer import LanguageModelTrainer, TextCorpus
import flair
import torch


flair.device = torch.device("cuda:1")


def train_flair_embeddings(
    corpus_path="/data/ubertext/for_flair",
    dictionary_path="/home/dima/Projects/flair_embeddings/flair_dictionary.pkl",
    lm_file="./language_model_backward_no_amp_accum_grad_fixed",
    is_forward_lm=True,
    hidden_size=1024,
    sequence_length=250,
    mini_batch_size=200,
    max_epochs=25,
    gpus=1,
):
    # load the default character dictionary
    dictionary: Dictionary = Dictionary.load(dictionary_path)

    # get your corpus, process forward and at the character level
    corpus = TextCorpus(corpus_path, dictionary, is_forward_lm, character_level=True)

    checkpoint = os.path.join(lm_file, "checkpoint.pt")

    if os.path.exists(checkpoint):
        trainer = LanguageModelTrainer.load_checkpoint(checkpoint, corpus)
    else:
        # train your language model
        language_model = LanguageModel(dictionary, is_forward_lm, hidden_size=hidden_size, nlayers=1)
        trainer = LanguageModelTrainer(language_model, corpus)

    trainer.train(
        lm_file,
        sequence_length=sequence_length,
        mini_batch_size=mini_batch_size,
        max_epochs=max_epochs,
        use_amp=False,
        checkpoint=True,
    )


train_flair_embeddings(is_forward_lm=False, mini_batch_size=480)