SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Track each row separately for stopping criteria (#29116)

Browse Source
This commit is contained in:
Raushan Turganbay
2024-02-26 21:06:16 +05:00
committed by GitHub
parent ece1b62b93
commit 8f2f0f0f85
3 changed files with 43 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
tests/generation/test_stopping_criteria.py
View File

@@ -54,37 +54,37 @@ class StoppingCriteriaTestCase(unittest.TestCase):
]
)
self.assertFalse(criteria(input_ids, scores))
self.assertFalse(all(criteria(input_ids, scores)))
input_ids, scores = self._get_tensors(9)
self.assertFalse(criteria(input_ids, scores))
self.assertFalse(all(criteria(input_ids, scores)))
input_ids, scores = self._get_tensors(10)
self.assertTrue(criteria(input_ids, scores))
self.assertTrue(all(criteria(input_ids, scores)))
def test_max_length_criteria(self):
criteria = MaxLengthCriteria(max_length=10)
input_ids, scores = self._get_tensors(5)
self.assertFalse(criteria(input_ids, scores))
self.assertFalse(all(criteria(input_ids, scores)))
input_ids, scores = self._get_tensors(9)
self.assertFalse(criteria(input_ids, scores))
self.assertFalse(all(criteria(input_ids, scores)))
input_ids, scores = self._get_tensors(10)
self.assertTrue(criteria(input_ids, scores))
self.assertTrue(all(criteria(input_ids, scores)))
def test_max_new_tokens_criteria(self):
criteria = MaxNewTokensCriteria(start_length=5, max_new_tokens=5)
input_ids, scores = self._get_tensors(5)
self.assertFalse(criteria(input_ids, scores))
self.assertFalse(all(criteria(input_ids, scores)))
input_ids, scores = self._get_tensors(9)
self.assertFalse(criteria(input_ids, scores))
self.assertFalse(all(criteria(input_ids, scores)))
input_ids, scores = self._get_tensors(10)
self.assertTrue(criteria(input_ids, scores))
self.assertTrue(all(criteria(input_ids, scores)))
criteria_list = StoppingCriteriaList([criteria])
self.assertEqual(criteria_list.max_length, 10)
@@ -93,10 +93,10 @@ class StoppingCriteriaTestCase(unittest.TestCase):
input_ids, scores = self._get_tensors(5)
criteria = MaxTimeCriteria(max_time=0.1)
self.assertFalse(criteria(input_ids, scores))
self.assertFalse(all(criteria(input_ids, scores)))
criteria = MaxTimeCriteria(max_time=0.1, initial_timestamp=time.time() - 0.2)
self.assertTrue(criteria(input_ids, scores))
self.assertTrue(all(criteria(input_ids, scores)))
def test_validate_stopping_criteria(self):
validate_stopping_criteria(StoppingCriteriaList([MaxLengthCriteria(10)]), 10)