diff --git a/concordia/utils/sampling_test.py b/concordia/utils/sampling_test.py new file mode 100644 index 000000000..654239f30 --- /dev/null +++ b/concordia/utils/sampling_test.py @@ -0,0 +1,66 @@ +# Copyright 2024 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the sampling helper functions.""" + +from absl.testing import absltest +from absl.testing import parameterized +from concordia.utils import sampling + + +class ExtractChoiceResponseTest(parameterized.TestCase): + + @parameterized.named_parameters( + ('single_letter', 'a', 'a'), + ('single_digit', '5', '5'), + ('letter_then_paren', 'a)', 'a'), + ('letter_then_period', 'a.', 'a'), + ('parenthesized_at_start', '(a)bar', 'a'), + ('parenthesized_in_middle', 'foo(a)bar', 'a'), + ('parenthesized_at_end', 'The answer is (b).', 'b'), + ('letter_then_paren_in_longer_string', 'a) because', 'a'), + ) + def test_extracts_choice(self, sample, expected): + self.assertEqual(sampling.extract_choice_response(sample), expected) + + @parameterized.named_parameters( + ('empty_string', ''), + ('no_parenthesized_choice', 'no choice here'), + ('word_without_closing_paren', '(abc'), + ) + def test_returns_none_when_no_choice(self, sample): + self.assertIsNone(sampling.extract_choice_response(sample)) + + +class DynamicallyAdjustTemperatureTest(parameterized.TestCase): + + @parameterized.named_parameters( + # max_attempts == 10, so the midpoint is 5.0. + ('first_attempt', 1, 10, 0.0), + ('second_attempt_below_midpoint', 2, 10, 0.5), + ('just_below_midpoint', 4, 10, 0.5), + ('exactly_at_midpoint', 5, 10, 0.0), + ('just_above_midpoint', 6, 10, 0.75), + ('final_attempt', 10, 10, 0.75), + ('zero_attempts', 0, 10, 0.0), + ) + def test_temperature(self, attempts, max_attempts, expected): + self.assertEqual( + sampling.dynamically_adjust_temperature(attempts, max_attempts), + expected, + ) + + +if __name__ == '__main__': + absltest.main() diff --git a/concordia/utils/text_test.py b/concordia/utils/text_test.py new file mode 100644 index 000000000..7fce497e5 --- /dev/null +++ b/concordia/utils/text_test.py @@ -0,0 +1,73 @@ +# Copyright 2023 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the string formatting utilities.""" + +from absl.testing import absltest +from absl.testing import parameterized +from concordia.utils import text + + +class WrapTest(parameterized.TestCase): + + def test_short_string_is_unchanged(self): + self.assertEqual(text.wrap('hello world', width=70), 'hello world') + + def test_empty_string(self): + self.assertEqual(text.wrap(''), '') + + def test_wraps_at_width(self): + self.assertEqual(text.wrap('aaaa bbbb cccc', width=5), 'aaaa\nbbbb\ncccc') + + def test_preserves_existing_newlines(self): + self.assertEqual(text.wrap('line one\nline two', width=70), + 'line one\nline two') + + def test_no_line_exceeds_width(self): + string = 'the quick brown fox jumps over the lazy dog ' * 5 + width = 20 + for line in text.wrap(string, width=width).split('\n'): + self.assertLessEqual(len(line), width) + + +class TruncateTest(parameterized.TestCase): + + def test_no_truncation_by_default(self): + self.assertEqual(text.truncate('hello world'), 'hello world') + + def test_truncates_to_max_length(self): + self.assertEqual(text.truncate('hello world', max_length=5), 'hello') + + def test_max_length_longer_than_string(self): + self.assertEqual(text.truncate('hello', max_length=100), 'hello') + + @parameterized.named_parameters( + ('single_delimiter', 'hello world', (' ',), 'hello'), + ('delimiter_absent', 'hello', (' ',), 'hello'), + ('multiple_delimiters', 'a.b,c', (',', '.'), 'a'), + ('stops_at_first_occurrence', 'a-b-c', ('-',), 'a'), + ) + def test_delimiters(self, string, delimiters, expected): + self.assertEqual( + text.truncate(string, delimiters=delimiters), expected) + + def test_max_length_and_delimiters_combined(self): + self.assertEqual( + text.truncate('hello world foo', max_length=11, delimiters=(' ',)), + 'hello', + ) + + +if __name__ == '__main__': + absltest.main()