Spaces:
Running on Zero
Running on Zero
fix: strip unused minicpm generation inputs
Browse files
tests/test_model_runtime.py
CHANGED
|
@@ -8,6 +8,7 @@ from hackathon_advisor.model_runtime import (
|
|
| 8 |
render_context,
|
| 9 |
runtime_status,
|
| 10 |
system_prompt,
|
|
|
|
| 11 |
)
|
| 12 |
from hackathon_advisor.zerogpu import gpu_task, zero_gpu_duration_seconds, zero_gpu_enabled
|
| 13 |
|
|
@@ -176,3 +177,11 @@ def test_zerogpu_duration_validates_positive_values(monkeypatch: pytest.MonkeyPa
|
|
| 176 |
monkeypatch.setenv("ADVISOR_ZERO_GPU_DURATION", "121")
|
| 177 |
with pytest.raises(RuntimeError, match="at most 120"):
|
| 178 |
zero_gpu_duration_seconds()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 8 |
render_context,
|
| 9 |
runtime_status,
|
| 10 |
system_prompt,
|
| 11 |
+
_strip_unused_generation_inputs,
|
| 12 |
)
|
| 13 |
from hackathon_advisor.zerogpu import gpu_task, zero_gpu_duration_seconds, zero_gpu_enabled
|
| 14 |
|
|
|
|
| 177 |
monkeypatch.setenv("ADVISOR_ZERO_GPU_DURATION", "121")
|
| 178 |
with pytest.raises(RuntimeError, match="at most 120"):
|
| 179 |
zero_gpu_duration_seconds()
|
| 180 |
+
|
| 181 |
+
|
| 182 |
+
def test_generation_inputs_drop_token_type_ids() -> None:
|
| 183 |
+
inputs = {"input_ids": [1], "attention_mask": [1], "token_type_ids": [0]}
|
| 184 |
+
|
| 185 |
+
_strip_unused_generation_inputs(inputs)
|
| 186 |
+
|
| 187 |
+
assert inputs == {"input_ids": [1], "attention_mask": [1]}
|