wehub-resource-sync
42 lines
1.6 KiB
Python
42 lines
1.6 KiB
Python
"""This file contains tests for the TextGeneration class."""
|
|
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
import shap
|
|
|
|
|
|
@pytest.mark.skipif(sys.platform == "win32", reason="Integer division bug in HuggingFace on Windows")
|
|
def test_call_function_text_generation():
|
|
"""Tests if target sentence from model and model wrapped in a function (mimics model agnostic scenario)
|
|
produces the same ids.
|
|
"""
|
|
torch = pytest.importorskip("torch")
|
|
transformers = pytest.importorskip("transformers")
|
|
|
|
name = "hf-internal-testing/tiny-random-BartModel"
|
|
tokenizer = transformers.AutoTokenizer.from_pretrained(name)
|
|
model = transformers.AutoModelForSeq2SeqLM.from_pretrained(name)
|
|
|
|
# Define function
|
|
def f(x):
|
|
inputs = tokenizer(x.tolist(), return_tensors="pt", padding=True)
|
|
with torch.no_grad():
|
|
out = model.generate(**inputs)
|
|
sentence = [tokenizer.decode(g, skip_special_tokens=True) for g in out]
|
|
return sentence
|
|
|
|
text_generation_for_pretrained_model = shap.models.TextGeneration(model, tokenizer=tokenizer, device="cpu")
|
|
text_generation_for_model_agnostic_scenario = shap.models.TextGeneration(f, device="cpu")
|
|
|
|
s = "This is a test statement for verifying text generation ids"
|
|
|
|
target_sentence_ids_for_pretrained_model = text_generation_for_pretrained_model(s)
|
|
target_sentence_for_pretrained_model = [
|
|
tokenizer.decode(g, skip_special_tokens=True) for g in target_sentence_ids_for_pretrained_model
|
|
]
|
|
target_sentence_for_model_agnostic_scenario = text_generation_for_model_agnostic_scenario(s)
|
|
|
|
assert target_sentence_for_pretrained_model[0] == target_sentence_for_model_agnostic_scenario[0]
|