wehub-resource-sync
61 lines
2.5 KiB
Python
61 lines
2.5 KiB
Python
import time
|
|
|
|
import pytest
|
|
|
|
|
|
def load_tokenizer_model(name: str, retries: int) -> tuple:
|
|
AutoTokenizer = pytest.importorskip("transformers").AutoTokenizer
|
|
AutoModelForSeq2SeqLM = pytest.importorskip("transformers").AutoModelForSeq2SeqLM
|
|
|
|
max_retries = retries # Use the parameter value
|
|
for attempt in range(max_retries):
|
|
try:
|
|
tokenizer = AutoTokenizer.from_pretrained(name)
|
|
model = AutoModelForSeq2SeqLM.from_pretrained(name)
|
|
return tokenizer, model
|
|
except OSError:
|
|
time.sleep(2**attempt) # Exponential backoff
|
|
raise OSError(f"Failed to load model and tokenizer after {max_retries} attempts")
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def basic_translation_scenario():
|
|
"""Create a basic transformers translation model and tokenizer."""
|
|
pytest.importorskip("torch")
|
|
# Use a *tiny* tokenizer model, to keep tests running as fast as possible.
|
|
# Nb. At time of writing, this pretrained model requires "protobuf==3.20.3".
|
|
# name = "mesolitica/finetune-translation-t5-super-super-tiny-standard-bahasa-cased"
|
|
# name = "Helsinki-NLP/opus-mt-en-es"
|
|
name = "hf-internal-testing/tiny-random-BartModel"
|
|
|
|
tokenizer, model = load_tokenizer_model(name=name, retries=5)
|
|
|
|
# define the input sentences we want to translate
|
|
data = [
|
|
"In this picture, there are four persons: my father, my mother, my brother and my sister.",
|
|
"Transformers have rapidly become the model of choice for NLP problems, replacing older recurrent neural network models",
|
|
]
|
|
|
|
return model, tokenizer, data
|
|
|
|
|
|
@pytest.fixture()
|
|
def causalml_synth_data():
|
|
"""
|
|
Generates synthetic data for causalml causal tree tests.
|
|
|
|
Unlike standard regression trees causal trees in causalml evaluate outcome conditioning on treatments
|
|
Thus, a causal tree estimates Y_hat|X,T=t, where t={0, 1,..., n}.
|
|
The simplest case is when T = {0, 1}, 0 - no treatment, 1 - some treatment.
|
|
"""
|
|
dataset = pytest.importorskip("causalml.dataset")
|
|
|
|
data_mode = 1 # Basic synthetic data mode with a difficult nuisance components and an easy treatment effect
|
|
sigma = 0.1 # Synthetic standard deviation of the error term
|
|
n_observations = 100 # The number of samples to generate
|
|
n_features = 8 # X in (Y_hat|X, T=0, Y_hat|X, T=1)
|
|
n_outcomes = 2 # Treatment conditioned outcomes: (Y_hat|X,T=0, Y_hat|X,T=1)
|
|
|
|
data = dataset.synthetic_data(mode=data_mode, n=n_observations, p=n_features, sigma=sigma)
|
|
return data, n_outcomes
|