Spaces:
Running on Zero
Running on Zero
File size: 2,185 Bytes
2f31638 0d0289f cad6462 0d0289f 2f31638 a71f0c7 2f31638 0d0289f cad6462 2f31638 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | """Regression tests for generation configuration."""
from __future__ import annotations
import unittest
from generation import (
ensure_bos_token,
gpu_duration_seconds,
head_tail_token_counts,
merge_eos_token_ids,
)
class GenerationConfigTests(unittest.TestCase):
def test_adds_missing_bos_token(self) -> None:
self.assertEqual(
ensure_bos_token("<|im_start|>user\nHello", "<|begin_of_text|>"),
"<|begin_of_text|><|im_start|>user\nHello",
)
def test_does_not_duplicate_bos_token(self) -> None:
prompt = "<|begin_of_text|><|im_start|>user\nHello"
self.assertEqual(ensure_bos_token(prompt, "<|begin_of_text|>"), prompt)
def test_preserves_all_model_terminators(self) -> None:
self.assertEqual(
merge_eos_token_ids([128001, 128008, 128009, 128256], 128256),
[128001, 128008, 128009, 128256],
)
def test_adds_tokenizer_fallback(self) -> None:
self.assertEqual(merge_eos_token_ids(128001, 128256), [128001, 128256])
def test_returns_none_without_valid_ids(self) -> None:
self.assertIsNone(merge_eos_token_ids(None, None))
def test_short_prompt_is_not_split(self) -> None:
self.assertEqual(head_tail_token_counts(100, 200, 50), (100, 0))
def test_large_prompt_preserves_prefix_and_maximizes_tail(self) -> None:
self.assertEqual(
head_tail_token_counts(30_000, 16_000, 4_096),
(4_096, 11_904),
)
def test_tiny_budget_still_preserves_latest_token(self) -> None:
self.assertEqual(head_tail_token_counts(100, 1, 4_096), (0, 1))
def test_gpu_duration_accounts_for_large_context_prefill(self) -> None:
self.assertEqual(
gpu_duration_seconds(90_000, 32, 16_384),
110,
)
def test_gpu_duration_keeps_small_requests_efficient(self) -> None:
self.assertEqual(gpu_duration_seconds(300, 32, 16_384), 30)
def test_gpu_duration_is_capped(self) -> None:
self.assertEqual(
gpu_duration_seconds(1_000_000, 1_024, 16_384),
120,
)
if __name__ == "__main__":
unittest.main()
|