|
import pytest |
|
|
|
from ding.model.template.language_transformer import LanguageTransformer |
|
|
|
|
|
@pytest.mark.unittest |
|
class TestNLPPretrainedModel: |
|
|
|
def check_model(self): |
|
test_pids = [1] |
|
cand_pids = [0, 2, 4] |
|
problems = [ |
|
"This is problem 0", "This is the first question", "Second problem is here", "Another problem", |
|
"This is the last problem" |
|
] |
|
ctxt_list = [problems[pid] for pid in test_pids] |
|
cands_list = [problems[pid] for pid in cand_pids] |
|
|
|
model = LanguageTransformer(model_name="bert-base-uncased", add_linear=True, embedding_size=256) |
|
scores = model(ctxt_list, cands_list) |
|
assert scores.shape == (1, 3) |
|
|
|
model = LanguageTransformer(model_name="bert-base-uncased", add_linear=False, embedding_size=256) |
|
scores = model(ctxt_list, cands_list) |
|
assert scores.shape == (1, 3) |
|
|
