diff --git a/integration-tests/fixtures/neuron/export_models.py b/integration-tests/fixtures/neuron/export_models.py index 836402ec..d4d0f01c 100644 --- a/integration-tests/fixtures/neuron/export_models.py +++ b/integration-tests/fixtures/neuron/export_models.py @@ -28,15 +28,6 @@ logger = logging.getLogger(__file__) # All model configurations below will be added to the neuron_model_config fixture MODEL_CONFIGURATIONS = { - "gpt2": { - "model_id": "gpt2", - "export_kwargs": { - "batch_size": 4, - "sequence_length": 1024, - "num_cores": 2, - "auto_cast_type": "fp16", - }, - }, "llama": { "model_id": "unsloth/Llama-3.2-1B-Instruct", "export_kwargs": { @@ -46,15 +37,6 @@ MODEL_CONFIGURATIONS = { "auto_cast_type": "fp16", }, }, - "mistral": { - "model_id": "optimum/mistral-1.1b-testing", - "export_kwargs": { - "batch_size": 4, - "sequence_length": 4096, - "num_cores": 2, - "auto_cast_type": "bf16", - }, - }, "qwen2": { "model_id": "Qwen/Qwen2.5-0.5B", "export_kwargs": { diff --git a/integration-tests/neuron/test_generate.py b/integration-tests/neuron/test_generate.py index f0804356..555b4eaa 100644 --- a/integration-tests/neuron/test_generate.py +++ b/integration-tests/neuron/test_generate.py @@ -20,9 +20,7 @@ async def test_model_single_request(tgi_service): ) assert response.details.generated_tokens == 17 greedy_expectations = { - "gpt2": "\n\nDeep learning is a new field of research that has been around for a while", "llama": " and How Does it Work?\nDeep learning is a subset of machine learning that uses artificial", - "mistral": "\nWhat is Deep Learning?\nDeep Learning is a type of machine learning that", "qwen2": " - Part 1\n\nDeep Learning is a subset of Machine Learning that is based on", "granite": "\n\nDeep Learning is a subset of Machine Learning, which is a branch of Art", } @@ -79,9 +77,7 @@ async def test_model_multiple_requests(tgi_service, neuron_generate_load): assert len(responses) == 4 expectations = { - "gpt2": "Deep learning is a new field of research that has been around for a while", "llama": "Deep learning is a subset of machine learning that uses artificial", - "mistral": "Deep Learning is a type of machine learning that", "qwen2": "Deep Learning is a subset of Machine Learning that is based on", "granite": "Deep Learning is a subset of Machine Learning, which is a branch of Art", }