wehub-resource-sync
91 lines
3.3 KiB
Python
91 lines
3.3 KiB
Python
"""This file contains tests for the TeacherForcingLogits class."""
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import shap
|
|
|
|
|
|
def test_falcon():
|
|
pytest.importorskip("torch")
|
|
transformers = pytest.importorskip("transformers")
|
|
requests = pytest.importorskip("requests")
|
|
name = "fxmarty/really-tiny-falcon-testing"
|
|
try:
|
|
tokenizer = transformers.AutoTokenizer.from_pretrained(name)
|
|
model = transformers.AutoModelForCausalLM.from_pretrained(
|
|
name, trust_remote_code=True, load_in_8bit=False, low_cpu_mem_usage=False
|
|
)
|
|
except requests.exceptions.RequestException:
|
|
pytest.xfail(reason="Connection error to transformers model")
|
|
|
|
model = model.eval()
|
|
|
|
s = ["I enjoy walking with my cute dog"]
|
|
gen_dict = dict(
|
|
max_new_tokens=100,
|
|
num_beams=5,
|
|
renormalize_logits=True,
|
|
no_repeat_ngram_size=8,
|
|
)
|
|
|
|
model.config.task_specific_params = dict()
|
|
model.config.task_specific_params["text-generation"] = gen_dict
|
|
shap_model = shap.models.TeacherForcing(model, tokenizer)
|
|
|
|
explainer = shap.Explainer(shap_model, tokenizer)
|
|
shap_values = explainer(s)
|
|
assert not np.isnan(np.sum(shap_values.values)) # type: ignore[union-attr]
|
|
|
|
|
|
def test_method_get_teacher_forced_logits_for_encoder_decoder_model():
|
|
"""Tests if get_teacher_forced_logits() works for encoder-decoder models."""
|
|
pytest.importorskip("torch")
|
|
transformers = pytest.importorskip("transformers")
|
|
requests = pytest.importorskip("requests")
|
|
|
|
name = "hf-internal-testing/tiny-random-BartModel"
|
|
try:
|
|
tokenizer = transformers.AutoTokenizer.from_pretrained(name)
|
|
model = transformers.AutoModelForSeq2SeqLM.from_pretrained(name)
|
|
except requests.exceptions.RequestException:
|
|
pytest.xfail(reason="Connection error to transformers model")
|
|
|
|
wrapped_model = shap.models.TeacherForcing(model, tokenizer, device="cpu")
|
|
|
|
source_sentence = np.array(
|
|
["This is a test statement for verifying working of teacher forcing logits functionality"]
|
|
)
|
|
target_sentence = np.array(["Testing teacher forcing logits functionality"])
|
|
|
|
# call the get teacher forced logits function
|
|
logits = wrapped_model.get_teacher_forced_logits(source_sentence, target_sentence)
|
|
|
|
assert not np.isnan(np.sum(logits))
|
|
|
|
|
|
def test_method_get_teacher_forced_logits_for_decoder_model():
|
|
"""Tests if get_teacher_forced_logits() works for decoder only models."""
|
|
pytest.importorskip("torch")
|
|
transformers = pytest.importorskip("transformers")
|
|
requests = pytest.importorskip("requests")
|
|
|
|
name = "hf-internal-testing/tiny-random-gpt2"
|
|
try:
|
|
tokenizer = transformers.AutoTokenizer.from_pretrained(name)
|
|
model = transformers.AutoModelForCausalLM.from_pretrained(name)
|
|
except requests.exceptions.RequestException:
|
|
pytest.xfail(reason="Connection error to transformers model")
|
|
|
|
model.config.is_decoder = True
|
|
|
|
wrapped_model = shap.models.TeacherForcing(model, tokenizer, device="cpu")
|
|
|
|
source_sentence = np.array(["This is a test statement for verifying"])
|
|
target_sentence = np.array(["working of teacher forcing logits functionality"])
|
|
|
|
# call the get teacher forced logits function
|
|
logits = wrapped_model.get_teacher_forced_logits(source_sentence, target_sentence)
|
|
|
|
assert not np.isnan(np.sum(logits))
|