[Bug]: Custom trained model not correctly loaded when using FlairEmbeddings

[petermartens1992]

Describe the bug

The F-scores shown after training a SequenceTagger model with FlairEmbeddings are way higher than when loading the model from file and doing an evaluation on it. Also the predictions from a loaded model seem way off compared to the initial F-scores.

The issue doesn't appear when initializing a new model and only loading the state from file:

sequence_tagger.load_state_dict(SequenceTagger.load(model_path).state_dict())

Also, when removing the FlairEmbeddings the F-scores between training and evaluation correctly align again.

See here the code to reproduce the issue (using flair 0.13.1):

import sys
import logging
import os

from typing import List
from pathlib import Path

from flair.datasets import UD_ENGLISH
from flair.data import Corpus
from flair.models import SequenceTagger
from flair.embeddings import TokenEmbeddings, WordEmbeddings, StackedEmbeddings, FlairEmbeddings, CharacterEmbeddings, OneHotEmbeddings, TransformerWordEmbeddings
from flair.datasets import ColumnCorpus
from flair.trainers import ModelTrainer

language = "nl"
model_type = "upos"

configs["train"] = {
    "language": language,
    "model_type": model_type,
    "tag_format": "BIOES",
    "rnn_hidden_size": 256,
    "rnn_layers": 1,
    "use_crf": True,
    "word_dropout": 0.2,
    "dropout": 0.2,
    "learning_rate": 0.1,
    "mini_batch_size": 32,
    "mini_batch_chunk_size": 32,
    "train_initial_hidden_state": True,
    "embeddings_storage_ratio": 1.0,
    "fine_tune_flair_embeddings": False,
}

output_model_name = configs["train"]["language"] + "-flair-" + configs["train"]["model_source"] + "-example"

output_path = "~/models/flair/" + configs["train"]["model_type"] + "/" + output_model_name
model_path = output_path + "/best-model.pt"
logger.info("model path: " + model_path)

# 1. what label do we want to predict?
label_type = configs["train"]["model_type"]

# 2. load the corpus
corpus = UD_ENGLISH().downsample(0.1)
print(corpus)

# 3. make the label dictionary from the corpus
label_dict = corpus.make_label_dictionary(label_type=label_type)
print(label_dict)

# 4. initialize embeddings
embedding_types: List[TokenEmbeddings] = [
    WordEmbeddings(configs["train"]["language"] + "-wiki"),

    CharacterEmbeddings(),

    FlairEmbeddings(configs["train"]["language"] + "-forward", fine_tune=configs["train"]["fine_tune_flair_embeddings"]),
    FlairEmbeddings(configs["train"]["language"] + "-backward", fine_tune=configs["train"]["fine_tune_flair_embeddings"])
]
embeddings: StackedEmbeddings = StackedEmbeddings(embeddings=embedding_types)

# 5. initialize sequence tagger
sequence_tagger = SequenceTagger(hidden_size=256,
                                 embeddings=embeddings,
                                 tag_dictionary=label_dict,
                                 tag_type=label_type)

# 6. initialize trainer
trainer = ModelTrainer(sequence_tagger, corpus)

# 7. start training
trainer.train(
    output_path,
    learning_rate=0.1,
    mini_batch_size=configs["train"]["mini_batch_size"],
    mini_batch_chunk_size=configs["train"]["mini_batch_chunk_size"],
    max_epochs=5
)

logger.info("evaluating model via load_state_dict...")
sequence_tagger.load_state_dict(SequenceTagger.load(model_path).state_dict())
logger.info(f'Model: "{sequence_tagger}"')

result = sequence_tagger.evaluate(
    data_points=corpus.test,
    gold_label_type=label_type,
    mini_batch_size=configs["train"]["mini_batch_size"],
    mini_batch_chunk_size=configs["train"]["mini_batch_chunk_size"],
    return_loss=False,
)

logger.info("detailed results:")
logger.info(result.detailed_results)
logger.info("main score:")
logger.info(str(result.main_score))
logger.info("classification report:")
logger.info(str(result.classification_report))
logger.info("scores:")
logger.info(str(result.scores))

logger.info("evaluating model via load...")
sequence_tagger = SequenceTagger.load(model_path)  # bad results when loading the model in the default way if the model contains flair embeddings
result = sequence_tagger.evaluate(
    data_points=corpus.test,
    gold_label_type=label_type,
    mini_batch_size=configs["train"]["mini_batch_size"],
    mini_batch_chunk_size=configs["train"]["mini_batch_chunk_size"],
    return_loss=False,
)

logger.info("detailed results:")
logger.info(result.detailed_results)
logger.info("main score:")
logger.info(str(result.main_score))
logger.info("classification report:")
logger.info(str(result.classification_report))
logger.info("scores:")
logger.info(str(result.scores))

To Reproduce

.

Expected behavior

.

Logs and Stack traces

No response

Screenshots

No response

Additional Context

No response

Environment

.