Use Qwen3-0.6B

Signed-off-by: SimJeg <[email protected]>

Author: SimJeg

tests/fixtures.py

@@ -14,38 +14,38 @@ def get_device():
 
 @pytest.fixture(scope="session")
 def unit_test_model():
-    model = AutoModelForCausalLM.from_pretrained("MaxJeblick/llama2-0b-unit-test").eval()
+    model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B").eval()
     return model.to(get_device())
 
 
 @pytest.fixture(scope="session")
 def unit_test_model_output_attention():
     model = AutoModelForCausalLM.from_pretrained(
-        "MaxJeblick/llama2-0b-unit-test", attn_implementation="eager"
+        "Qwen/Qwen3-0.6B", attn_implementation="eager"
     ).eval()
     return model.to(get_device())
 
 
 @pytest.fixture(scope="session")
-def danube_500m_model():
-    model = AutoModelForCausalLM.from_pretrained("h2oai/h2o-danube3-500m-chat").eval()
+def qwen3_600m_model():
+    model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B").eval()
     return model.to(get_device())
 
 
 @pytest.fixture(scope="session")
 def kv_press_unit_test_pipeline():
     return pipeline(
         "kv-press-text-generation",
-        model="maxjeblick/llama2-0b-unit-test",
+        model="Qwen/Qwen3-0.6B",
         device=get_device(),
     )
 
 
 @pytest.fixture(scope="session")
-def kv_press_danube_pipeline():
+def kv_press_qwen3_600m_pipeline():
     return pipeline(
         "kv-press-text-generation",
-        model="h2oai/h2o-danube3-500m-chat",
+        model="Qwen/Qwen3-0.6B",
         device=get_device(),
     )
 

tests/test_decoding_compression.py

@@ -31,7 +31,7 @@ def test_decoding_compression(token_buffer_size):
     """Test that DecodingPress compresses the cache during decoding."""
 
     # Initialize pipeline with a small model
-    pipe = pipeline("kv-press-text-generation", model="MaxJeblick/llama2-0b-unit-test", device_map="auto")
+    pipe = pipeline("kv-press-text-generation", model="Qwen/Qwen3-0.6B", device_map="auto")
 
     # Create a DecodingPress with KnormPress
     press = DecodingPress(
@@ -65,7 +65,7 @@ def test_prefill_decoding_press_calls_both_phases():
     """Test that PrefillDecodingPress calls both prefilling and decoding presses."""
 
     # Initialize pipeline
-    pipe = pipeline("kv-press-text-generation", model="MaxJeblick/llama2-0b-unit-test", device_map="auto")
+    pipe = pipeline("kv-press-text-generation", model="Qwen/Qwen3-0.6B", device_map="auto")
 
     # Create PrefillDecodingPress with both presses
     combined_press = PrefillDecodingPress(
@@ -99,7 +99,7 @@ def test_decoding_press_without_prefill():
     """Test that DecodingPress works correctly when used standalone (no prefill compression)."""
 
     # Initialize pipeline
-    pipe = pipeline("kv-press-text-generation", model="MaxJeblick/llama2-0b-unit-test", device_map="auto")
+    pipe = pipeline("kv-press-text-generation", model="Qwen/Qwen3-0.6B", device_map="auto")
 
     # Create DecodingPress only
     decoding_press = DecodingPress(base_press=KnormPress(compression_ratio=0.4), compression_interval=5, target_size=64)
@@ -129,7 +129,7 @@ def test_prefill_decoding_press_decoding_only():
     """Test PrefillDecodingPress with only decoding press (no prefill compression)."""
 
     # Initialize pipeline
-    pipe = pipeline("kv-press-text-generation", model="MaxJeblick/llama2-0b-unit-test", device_map="auto")
+    pipe = pipeline("kv-press-text-generation", model="Qwen/Qwen3-0.6B", device_map="auto")
 
     # Create PrefillDecodingPress with only decoding press
     combined_press = PrefillDecodingPress(
@@ -167,7 +167,7 @@ def test_decoding_press_equivalence():
     torch.manual_seed(42)
 
     # Initialize pipeline
-    pipe = pipeline("kv-press-text-generation", model="MaxJeblick/llama2-0b-unit-test", device_map="auto")
+    pipe = pipeline("kv-press-text-generation", model="Qwen/Qwen3-0.6B", device_map="auto")
 
     # Create standalone decoding press
     decoding_press = DecodingPress(base_press=KnormPress(compression_ratio=0.5), compression_interval=3, target_size=52)
@@ -222,7 +222,7 @@ def test_all_presses_work_with_decoding_press(press_config):
     """Test that all default presses work as base presses for DecodingPress."""
 
     # Initialize pipeline
-    pipe = pipeline("kv-press-text-generation", model="MaxJeblick/llama2-0b-unit-test", device_map="auto")
+    pipe = pipeline("kv-press-text-generation", model="Qwen/Qwen3-0.6B", device_map="auto")
 
     # Get press class and use the first (easier) configuration
     press_cls = press_config["cls"]
@@ -274,7 +274,7 @@ def test_all_presses_work_with_decoding_press(press_config):
 def test_compression_actually_reduces_memory():
     """Test that compression actually reduces memory usage compared to no compression."""
 
-    pipe = pipeline("kv-press-text-generation", model="MaxJeblick/llama2-0b-unit-test", device_map="auto")
+    pipe = pipeline("kv-press-text-generation", model="Qwen/Qwen3-0.6B", device_map="auto")
 
     context = "The quick brown fox jumps over the lazy dog. " * 15  # Long context
     question = "What animal jumps over the dog?"

tests/test_pipeline.py

@@ -11,8 +11,8 @@
 
 from kvpress import ExpectedAttentionPress
 from kvpress.pipeline import KVPressTextGenerationPipeline
-from tests.fixtures import danube_500m_model  # noqa: F401
-from tests.fixtures import kv_press_danube_pipeline  # noqa: F401
+from tests.fixtures import qwen3_600m_model  # noqa: F401
+from tests.fixtures import kv_press_qwen3_600m_pipeline  # noqa: F401
 from tests.fixtures import unit_test_model  # noqa: F401
 from tests.fixtures import kv_press_llama3_2_flash_attn_pipeline, kv_press_unit_test_pipeline  # noqa: F401
 
@@ -94,9 +94,9 @@ def test_pipeline_no_press_works(kv_press_unit_test_pipeline, caplog):  # noqa:
     kv_press_unit_test_pipeline(context, question=question)
 
 
-def test_pipeline_answer_is_correct(danube_500m_model, caplog):  # noqa: F811
+def test_pipeline_answer_is_correct(qwen3_600m_model, caplog):  # noqa: F811
     with caplog.at_level(logging.DEBUG):
-        answers = generate_answer(danube_500m_model)
+        answers = generate_answer(qwen3_600m_model)
 
     for answer in answers:
         assert answer == "This article was written on January 1, 2022."
@@ -107,13 +107,13 @@ def test_pipeline_answer_is_correct(danube_500m_model, caplog):  # noqa: F811
 
 
 @pytest.mark.skipif(not is_optimum_quanto_available(), reason="Optimum Quanto is not available")
-def test_pipeline_with_quantized_cache(kv_press_danube_pipeline, caplog):  # noqa: F811
+def test_pipeline_with_quantized_cache(kv_press_qwen3_600m_pipeline, caplog):  # noqa: F811
     with caplog.at_level(logging.DEBUG):
         context = "This is a test article. It was written on 2022-01-01."
         questions = ["When was this article written?"]
         press = ExpectedAttentionPress(compression_ratio=0.4)
-        cache = QuantoQuantizedCache(config=kv_press_danube_pipeline.model.config, nbits=4)
-        answers = kv_press_danube_pipeline(context, questions=questions, press=press, cache=cache)["answers"]
+        cache = QuantoQuantizedCache(config=kv_press_qwen3_600m_pipeline.model.config, nbits=4)
+        answers = kv_press_qwen3_600m_pipeline(context, questions=questions, press=press, cache=cache)["answers"]
 
     assert len(answers) == 1
     assert isinstance(answers[0], str)