JacobLinCool commited on
Commit
4e530e4
·
verified ·
1 Parent(s): e5234dc

fix: strip unused minicpm generation inputs

Browse files
Files changed (1) hide show
  1. tests/test_model_runtime.py +9 -0
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]}