hf_text-generation-inference/server/tests/conftest.py

21 lines
453 B
Python
Raw Normal View History

2022-12-08 10:49:33 -07:00
import pytest
2023-03-07 10:52:22 -07:00
from text_generation_server.pb import generate_pb2
2022-12-08 10:49:33 -07:00
@pytest.fixture
def default_pb_parameters():
2022-12-15 09:03:56 -07:00
return generate_pb2.NextTokenChooserParameters(
2022-12-08 10:49:33 -07:00
temperature=1.0,
repetition_penalty=1.0,
2022-12-08 10:49:33 -07:00
top_k=0,
top_p=1.0,
typical_p=1.0,
2022-12-08 10:49:33 -07:00
do_sample=False,
)
2022-12-12 10:25:22 -07:00
@pytest.fixture
def default_pb_stop_parameters():
return generate_pb2.StoppingCriteriaParameters(stop_sequences=[], max_new_tokens=10)