From 31296e5005c849443f2155df1f9e6a9d1a264d13 Mon Sep 17 00:00:00 2001 From: KexinFeng Date: Thu, 6 Jul 2023 10:45:24 -0700 Subject: [PATCH] unit test --- .../setup/djl_python/scheduler/seq_batch_scheduler.py | 3 +-- engines/python/setup/djl_python/tests/test_scheduler.py | 6 ------ 2 files changed, 1 insertion(+), 8 deletions(-) diff --git a/engines/python/setup/djl_python/scheduler/seq_batch_scheduler.py b/engines/python/setup/djl_python/scheduler/seq_batch_scheduler.py index 416f499c4f..ab9fdf0d01 100644 --- a/engines/python/setup/djl_python/scheduler/seq_batch_scheduler.py +++ b/engines/python/setup/djl_python/scheduler/seq_batch_scheduler.py @@ -43,8 +43,7 @@ def __init__(self, lm_block: LMBlock, default_search_algorithm: str, self.results: Dict[int, List[int]] = defaultdict(list) self.seq_batchers: Dict[ - Type[SeqBatcher]:List[SeqBatcher]] = defaultdict( - list) # {key: List[SeqBatcher]} + Type[SeqBatcher]:List[SeqBatcher]] = defaultdict(list) self.lru_kv_cache = OrderedDict() self.lru_max_size = 10 diff --git a/engines/python/setup/djl_python/tests/test_scheduler.py b/engines/python/setup/djl_python/tests/test_scheduler.py index 7b9aa5a037..fe757d0d57 100644 --- a/engines/python/setup/djl_python/tests/test_scheduler.py +++ b/engines/python/setup/djl_python/tests/test_scheduler.py @@ -161,10 +161,6 @@ def test_sampling_scheduler(self): results = scheduler.collect_results() - print('=====++++++++++==============') - for i, ret in results.items(): - print('\n{}:'.format(i), tokenizer.decode(ret[:30])) - assert tokenizer.decode(results[0][:30]) == "Memories follow me left and right. I can't remember the last " \ "time I had to do this, and I'm glad that I did.\n" assert tokenizer.decode( @@ -174,8 +170,6 @@ def test_sampling_scheduler(self): results[2][:30] ) == "There's a time that I remember, when I did not know what to do with my life. I was in a very bad mood. I was" - print('======+++++++++++++=========') - def test_contrastive_scheduler(self): model_id = "gpt2" model = GPT2LMHeadModel.from_pretrained(model_id)