import re
from enum import Enum
from pydantic import BaseModel
import pytest
import torch
import transformers
import outlines
from outlines.inputs import Chat
from outlines.models.transformers import (
Transformers,
TransformerTokenizer,
TransformersTypeAdapter,
)
from outlines.types import Regex
TEST_MODEL = "erwanf/gpt2-mini"
TEST_MODEL_MAMBA = "hf-internal-testing/tiny-random-MambaForCausalLM"
TEST_MODEL_BART = "trl-internal-testing/tiny-BartModel"
def test_transformers_instantiate_invalid():
with pytest.raises(ValueError):
outlines.from_transformers(
transformers.AutoModelForCausalLM.from_pretrained(TEST_MODEL),
int,
)
def test_transformers_instantiate_simple():
model = outlines.from_transformers(
transformers.AutoModelForCausalLM.from_pretrained(TEST_MODEL),
transformers.AutoTokenizer.from_pretrained(TEST_MODEL),
)
assert isinstance(model, Transformers)
assert isinstance(model.tokenizer, TransformerTokenizer)
assert isinstance(model.type_adapter, TransformersTypeAdapter)
assert model.tensor_library_name == "torch"
def test_transformers_instantiate_mamba():
model = outlines.from_transformers(
transformers.MambaForCausalLM.from_pretrained(TEST_MODEL_MAMBA),
transformers.AutoTokenizer.from_pretrained(TEST_MODEL),
)
assert isinstance(model, Transformers)
def test_transformers_instantiate_tokenizer_kwargs_dtype():
model = outlines.from_transformers(
transformers.AutoModelForCausalLM.from_pretrained(TEST_MODEL),
transformers.AutoTokenizer.from_pretrained(
TEST_MODEL, additional_special_tokens=["<t1>", "<t2>"]
),
device_dtype=torch.bfloat16,
)
assert "<t1>" in model.tokenizer.special_tokens
assert "<t2>" in model.tokenizer.special_tokens
assert model.device_dtype == torch.bfloat16
@pytest.fixture
def model():
model = outlines.from_transformers(
transformers.AutoModelForCausalLM.from_pretrained(TEST_MODEL),
transformers.AutoTokenizer.from_pretrained(TEST_MODEL),
)
chat_template = '{% for message in messages %}{{ message.role }}: {{ message.content }}{% endfor %}'
model.type_adapter.tokenizer.chat_template = chat_template
return model
@pytest.fixture
def model_bart():
model = outlines.from_transformers(
transformers.BartForConditionalGeneration.from_pretrained(TEST_MODEL_BART),
transformers.BartTokenizer.from_pretrained(TEST_MODEL_BART),
)
return model
def test_transformers_simple(model):
result = model.generate("Respond with one word. Not more.", None)
assert isinstance(result, str)
def test_transformers_call(model, model_bart):
result = model("Respond with one word. Not more.")
assert isinstance(result, str)
model.device_dtype = torch.bfloat16
result = model("Respond with one word. Not more.")
assert isinstance(result, str)
result = model_bart("Respond with one word. Not more.")
assert isinstance(result, str)
def test_transformers_chat(model):
result = model(
Chat(messages=[
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "What is the capital of France?"},
])
)
assert isinstance(result, str)
def test_transformers_inference_kwargs(model):
result = model("Respond with one word. Not more.", max_new_tokens=100)
assert isinstance(result, str)
def test_transformers_invalid_inference_kwargs(model):
with pytest.raises(ValueError):
model("Respond with one word. Not more.", foo="bar")
def test_transformers_regex(model):
result = model("Give a number between 0 and 9.", Regex(r"[0-9]"))
assert isinstance(result, str)
assert re.match(r"[0-9]", result)
def test_transformers_json(model):
class Character(BaseModel):
name: str
result = model("Create a character with a name.", Character)
assert "name" in result
def test_transformers_choice(model):
class Foo(Enum):
cat = "cat"
dog = "dog"
result = model("Cat or dog?", Foo)
assert result in ["cat", "dog"]
def test_transformers_multiple_samples(model):
result = model("Respond with one word. Not more.")
assert isinstance(result, str)
result = model(
"Respond with one word. Not more.", num_return_sequences=2, do_sample=True
)
assert isinstance(result, list)
assert len(result) == 2
def test_transformers_batch(model):
result = model.batch(
["Respond with one word. Not more.", "Respond with one word. Not more."]
)
assert isinstance(result, list)
assert len(result) == 2
result = model.batch(
["Respond with one word. Not more.", "Respond with one word. Not more."],
num_return_sequences=2,
do_sample=True,
)
assert isinstance(result, list)
assert len(result) == 2
for item in result:
assert isinstance(item, list)
assert len(item) == 2
result = model.batch(
[
Chat(messages=[
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "What is the capital of France?"},
]),
Chat(messages=[
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "What is the capital of France?"},
]),
],
)
assert isinstance(result, list)
assert len(result) == 2
def test_transformers_multiple_samples_constrained(model):
class Foo(Enum):
cat = "cat"
dog = "dog"
result = model("Cat or dog?", Foo, num_return_sequences=2, do_sample=True)
assert isinstance(result, list)
assert len(result) == 2
assert result[0] in ["cat", "dog"]
assert result[1] in ["cat", "dog"]
def test_transformers_batch_constrained(model):
class Foo(Enum):
cat = "cat"
dog = "dog"
result = model.batch(
["Cat or dog?", "Cat or dog?"],
Foo,
)
assert isinstance(result, list)
assert len(result) == 2
assert result[0] in ["cat", "dog"]
assert result[1] in ["cat", "dog"]
result = model.batch(
["Cat or dog?", "Cat or dog?"],
Foo,
num_return_sequences=2,
do_sample=True,
)
assert isinstance(result, list)
assert len(result) == 2
for item in result:
assert isinstance(item, list)
assert len(item) == 2
assert item[0] in ["cat", "dog"]
assert item[1] in ["cat", "dog"]
def test_transformers_streaming(model):
with pytest.raises(NotImplementedError, match="Streaming is not implemented"):
model.stream("Respond with one word. Not more.")
@pytest.mark.parametrize(
"model_name",
[
TEST_MODEL,
"HuggingFaceTB/SmolLM2-135M"
],
)
def test_transformers_parametrized_smoke(model_name):
"""
Smoke test to ensure basic constrained generation works across
different tokenizers.
"""
hf_model = transformers.AutoModelForCausalLM.from_pretrained(model_name)
hf_model.eval()
hf_tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
model = outlines.from_transformers(hf_model, hf_tokenizer)
prompt = "Is 1+1=2? Answer Yes or No:"
constraint = Regex(r"\s*(Yes|No)")
out = model(
prompt,
constraint,
max_new_tokens=5,
do_sample=False,
)
assert out.strip() in {"Yes", "No"}