-
Notifications
You must be signed in to change notification settings - Fork 1.8k
/
test_modeling_prediction_head.py
29 lines (22 loc) · 987 Bytes
/
test_modeling_prediction_head.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
import logging
from haystack.modeling.model.adaptive_model import AdaptiveModel
from haystack.modeling.model.language_model import LanguageModel
from haystack.modeling.model.prediction_head import QuestionAnsweringHead
from haystack.modeling.utils import set_all_seeds, initialize_device_settings
def test_prediction_head_load_save(tmp_path, caplog=None):
if caplog:
caplog.set_level(logging.CRITICAL)
set_all_seeds(seed=42)
devices, n_gpu = initialize_device_settings(use_cuda=False)
lang_model = "bert-base-german-cased"
language_model = LanguageModel.load(lang_model)
prediction_head = QuestionAnsweringHead()
model = AdaptiveModel(
language_model=language_model,
prediction_heads=[prediction_head],
embeds_dropout_prob=0.1,
lm_output_types=["per_sequence"],
device=devices[0])
model.save(tmp_path)
model_loaded = AdaptiveModel.load(tmp_path, device='cpu')
assert model_loaded is not None