Add stop to lf.LMSamplingOptions and enable it for OpenAI models.

PiperOrigin-RevId: 591402329
google · Dec 16, 2023 · 3b90e92 · 3b90e92
1 parent 0c68bf7
commit 3b90e92
Show file tree

Hide file tree

Showing 5 changed files with 23 additions and 12 deletions.
diff --git a/langfun/core/eval/base_test.py b/langfun/core/eval/base_test.py
@@ -101,7 +101,7 @@ def test_basics(self):
     self.assertEqual(s.dir, os.path.join(s.root_dir, s.id))
     self.assertEqual(s.hash, s.clone().hash)
     # Test persistent hash.
-    self.assertEqual(s.hash, '0c857e07')
+    self.assertEqual(s.hash, 'c76d4fe6')
     self.assertEqual(
         s.hash, s.clone(override={'max_workers': 2, 'lm.timeout': 20}).hash
     )
@@ -323,7 +323,7 @@ def test_search_space(self):
         s.children[0].dir, os.path.join(s.root_dir, s.children[0].id)
     )
     # Test persistent hash.
-    self.assertEqual(s.hash, 'cbb0adcf')
+    self.assertEqual(s.hash, 'e987475a')
 
     summary = s.run(verbose=True)
     self.assertEqual(len(summary.evaluations), 2)
@@ -451,7 +451,7 @@ def test_run(self):
         ],
     )
     # Test for persistent hash.
-    self.assertEqual(s.hash, 'ec4758d3')
+    self.assertEqual(s.hash, 'bb86a963')
     s.run()
     expected = {
         s.children[0].id: dict(

diff --git a/langfun/core/langfunc_test.py b/langfun/core/langfunc_test.py
@@ -73,11 +73,11 @@ def test_call(self):
     print(repr(l))
     self.assertEqual(
         repr(l),
-        "LangFunc(template_str='Hello', clean=True, "
-        'lm=ExcitedEchoer(sampling_options=LMSamplingOptions(temperature=0.0, '
-        'max_tokens=1024, n=1, top_k=40, top_p=None, random_seed=None), '
-        'cache=None, timeout=120.0, max_attempts=5, retry_interval=(5, 60), '
-        'exponential_backoff=True, debug=False))',
+        "LangFunc(template_str='Hello', clean=True,"
+        ' lm=ExcitedEchoer(sampling_options=LMSamplingOptions(temperature=0.0,'
+        ' max_tokens=1024, n=1, top_k=40, top_p=None, stop=None,'
+        ' random_seed=None), cache=None, timeout=120.0, max_attempts=5,'
+        ' retry_interval=(5, 60), exponential_backoff=True, debug=False))',
     )
 
     l = LangFunc('Hello')

diff --git a/langfun/core/language_model.py b/langfun/core/language_model.py
@@ -80,6 +80,15 @@ class LMSamplingOptions(component.Component):
           '`top_p` but not both.'
       ),
   ] = None
+  stop: Annotated[
+      list[str] | None,
+      (
+          'A list of stop sequences that prevent LLMs from outputting '
+          'more tokens. For example, when `stop` is set to ["User:", "Model:"] '
+          'LLMs will stop to emit more tokens when `User:` or '
+          '`Model:` is reached.'
+      ),
+  ] = None
   random_seed: Annotated[
       int | None, 'A fixed random seed used during model inference.'
   ] = None

diff --git a/langfun/core/llms/openai.py b/langfun/core/llms/openai.py
@@ -170,6 +170,8 @@ def _get_request_args(
 
     if options.top_p is not None:
       args['top_p'] = options.top_p
+    if options.stop:
+      args['stop'] = options.stop
     return args
 
   def _sample(self, prompts: list[lf.Message]) -> list[LMSamplingResult]:

diff --git a/langfun/core/llms/openai_test.py b/langfun/core/llms/openai_test.py
@@ -109,17 +109,17 @@ def test_get_request_args(self):
     )
     self.assertEqual(
         openai.Gpt4(api_key='test_key')._get_request_args(
-            lf.LMSamplingOptions(
-                temperature=1.0,
-                n=1)),
+            lf.LMSamplingOptions(temperature=1.0, stop=['\n'], n=1)
+        ),
         dict(
             model='gpt-4',
             n=1,
             temperature=1.0,
             max_tokens=1024,
             stream=False,
             timeout=120.0,
-        )
+            stop=['\n'],
+        ),
     )
 
   def test_call_completion(self):