forked from flairNLP/flair
/
test_text_regressor.py
56 lines (35 loc) · 1.69 KB
/
test_text_regressor.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
import pytest
from typing import Tuple
from flair.data import Dictionary, Corpus
from flair.data_fetcher import NLPTaskDataFetcher, NLPTask
from flair.embeddings import WordEmbeddings, DocumentRNNEmbeddings
from flair.models.text_regression_model import TextRegressor
# from flair.trainers.trainer_regression import RegressorTrainer
from flair.trainers import ModelTrainer
def init(tasks_base_path) -> Tuple[Corpus, TextRegressor, ModelTrainer]:
corpus = NLPTaskDataFetcher.load_corpus(NLPTask.REGRESSION, tasks_base_path)
glove_embedding: WordEmbeddings = WordEmbeddings("glove")
document_embeddings: DocumentRNNEmbeddings = DocumentRNNEmbeddings(
[glove_embedding], 128, 1, False, 64, False, False
)
model = TextRegressor(document_embeddings)
trainer = ModelTrainer(model, corpus)
return corpus, model, trainer
def test_labels_to_indices(tasks_base_path):
corpus, model, trainer = init(tasks_base_path)
result = model._labels_to_indices(corpus.train)
for i in range(len(corpus.train)):
expected = round(float(corpus.train[i].labels[0].value), 3)
actual = round(float(result[i].item()), 3)
assert expected == actual
def test_trainer_evaluation(tasks_base_path):
corpus, model, trainer = init(tasks_base_path)
expected = model.evaluate(corpus.dev)
assert expected is not None
# def test_trainer_results(tasks_base_path):
# corpus, model, trainer = init(tasks_base_path)
# results = trainer.train("regression_train/", max_epochs=1)
# assert results["test_score"] > 0
# assert len(results["dev_loss_history"]) == 1
# assert len(results["dev_score_history"]) == 1
# assert len(results["train_loss_history"]) == 1