Fix Santacoder test.

This commit is contained in:
Nicolas Patry 2024-07-02 16:55:48 +02:00
parent ce913b874b
commit db9acc4418
No known key found for this signature in database
GPG Key ID: E939E8CC91A1C674

View File

@ -1,13 +1,12 @@
import pytest
from text_generation_server.pb import generate_pb2
from text_generation_server.models.causal_lm import CausalLMBatch
from text_generation_server.models.santacoder import SantaCoder
from text_generation_server.models.causal_lm import CausalLMBatch, CausalLM
@pytest.fixture(scope="session")
def default_santacoder():
return SantaCoder("bigcode/santacoder")
return CausalLM.fallback(model_id="bigcode/santacoder")
@pytest.fixture