langfun 0.0.2.dev20240414__py3-none-any.whl → 0.0.2.dev20240418__py3-none-any.whl

langfun/__init__.py

@@ -34,6 +34,7 @@ score = structured.score
 generate_class = structured.generate_class
 
 source_form = structured.source_form
+function_gen = structured.function_gen
 
 from langfun.core import eval  # pylint: disable=redefined-builtin
 from langfun.core import templates

langfun/core/__init__.py

@@ -99,6 +99,7 @@ from langfun.core.modality import ModalityRef
 from langfun.core.language_model import LanguageModel
 from langfun.core.language_model import LMSample
 from langfun.core.language_model import LMSamplingOptions
+from langfun.core.language_model import LMSamplingUsage
 from langfun.core.language_model import LMSamplingResult
 from langfun.core.language_model import LMScoringResult
 from langfun.core.language_model import LMCache

langfun/core/eval/base.py

@@ -1565,6 +1565,7 @@ class Summary(pg.Object):
         results.append(
             pg.Dict(
                 experiment=entry,
+                dir=entry.dir,
                 metrics=entry.result.metrics if entry.result else None,
             )
         )

langfun/core/eval/base_test.py

@@ -194,6 +194,7 @@ class EvaluationTest(unittest.TestCase):
             cache_seed=0,
             score=1.0,
             logprobs=None,
+            usage=lf.LMSamplingUsage(387, 24, 411),
             tags=['lm-response', 'lm-output', 'transformed'],
         ),
     )

langfun/core/langfunc_test.py

@@ -82,7 +82,9 @@ class LangFuncCallTest(unittest.TestCase):
     self.assertEqual(i.tags, ['rendered'])
 
     r = l()
-    self.assertEqual(r, message.AIMessage('Hello!!!', score=0.0, logprobs=None))
+    self.assertEqual(
+        r, message.AIMessage('Hello!!!', score=0.0, logprobs=None, usage=None)
+    )
     self.assertEqual(r.tags, ['lm-response', 'lm-output'])
     self.assertEqual(r.source, message.UserMessage('Hello'))
     self.assertEqual(r.source.tags, ['rendered', 'lm-input'])
@@ -106,7 +108,7 @@ class LangFuncCallTest(unittest.TestCase):
       self.assertEqual(l.render(), 'Hello')
       r = l()
       self.assertEqual(
-          r, message.AIMessage('Hello!!!', score=0.0, logprobs=None)
+          r, message.AIMessage('Hello!!!', score=0.0, logprobs=None, usage=None)
       )
       self.assertEqual(r.tags, ['lm-response', 'lm-output'])
 

langfun/core/language_model.py

@@ -47,6 +47,14 @@ class LMSample(pg.Object):
   ] = None
 
 
+class LMSamplingUsage(pg.Object):
+  """Usage information per completion."""
+
+  prompt_tokens: int
+  completion_tokens: int
+  total_tokens: int
+
+
 class LMSamplingResult(pg.Object):
   """Language model response."""
 
@@ -58,6 +66,11 @@ class LMSamplingResult(pg.Object):
       ),
   ] = []
 
+  usage: Annotated[
+      LMSamplingUsage | None,
+      'Usage information. Currently only OpenAI models are supported.',
+  ] = None
+
 
 class LMSamplingOptions(component.Component):
   """Language model sampling options."""
@@ -424,6 +437,8 @@ class LanguageModel(component.Component):
       logprobs = result.samples[0].logprobs
       response.set('score', result.samples[0].score)
       response.metadata.logprobs = logprobs
+      response.metadata.usage = result.usage
+
       elapse = time.time() - request_start
       self._debug(prompt, response, call_counter, elapse)
       return response

langfun/core/language_model_test.py

@@ -38,9 +38,19 @@ class MockModel(lm_lib.LanguageModel):
     def fake_sample(prompts):
       if context.attempt >= self.failures_before_attempt:
         return [
-            lm_lib.LMSamplingResult([lm_lib.LMSample(  # pylint: disable=g-complex-comprehension
-                response=prompt.text * self.sampling_options.top_k,
-                score=self.sampling_options.temperature or -1.0)])
+            lm_lib.LMSamplingResult(
+                [
+                    lm_lib.LMSample(  # pylint: disable=g-complex-comprehension
+                        response=prompt.text * self.sampling_options.top_k,
+                        score=self.sampling_options.temperature or -1.0,
+                    )
+                ],
+                usage=lm_lib.LMSamplingUsage(
+                    prompt_tokens=100,
+                    completion_tokens=100,
+                    total_tokens=200,
+                ),
+            )
             for prompt in prompts
         ]
       context.attempt += 1
@@ -100,8 +110,14 @@ class LanguageModelTest(unittest.TestCase):
     self.assertEqual(
         lm.sample(prompts=['foo', 'bar']),
         [
-            lm_lib.LMSamplingResult([lm_lib.LMSample('foo', score=-1.0)]),
-            lm_lib.LMSamplingResult([lm_lib.LMSample('bar', score=-1.0)]),
+            lm_lib.LMSamplingResult(
+                [lm_lib.LMSample('foo', score=-1.0)],
+                usage=lm_lib.LMSamplingUsage(100, 100, 200),
+            ),
+            lm_lib.LMSamplingResult(
+                [lm_lib.LMSample('bar', score=-1.0)],
+                usage=lm_lib.LMSamplingUsage(100, 100, 200),
+            ),
         ],
     )
     # Test override sampling_options.
@@ -112,10 +128,12 @@ class LanguageModelTest(unittest.TestCase):
         ),
         [
             lm_lib.LMSamplingResult(
-                [lm_lib.LMSample('foo' * 2, score=0.5)]
+                [lm_lib.LMSample('foo' * 2, score=0.5)],
+                usage=lm_lib.LMSamplingUsage(100, 100, 200),
             ),
             lm_lib.LMSamplingResult(
-                [lm_lib.LMSample('bar' * 2, score=0.5)]
+                [lm_lib.LMSample('bar' * 2, score=0.5)],
+                usage=lm_lib.LMSamplingUsage(100, 100, 200),
             ),
         ],
     )
@@ -123,18 +141,26 @@ class LanguageModelTest(unittest.TestCase):
     self.assertEqual(
         lm.sample(prompts=['foo', 'bar'], temperature=1.0),
         [
-            lm_lib.LMSamplingResult([lm_lib.LMSample('foo', score=1.0)]),
-            lm_lib.LMSamplingResult([lm_lib.LMSample('bar', score=1.0)]),
+            lm_lib.LMSamplingResult(
+                [lm_lib.LMSample('foo', score=1.0)],
+                usage=lm_lib.LMSamplingUsage(100, 100, 200),
+            ),
+            lm_lib.LMSamplingResult(
+                [lm_lib.LMSample('bar', score=1.0)],
+                usage=lm_lib.LMSamplingUsage(100, 100, 200),
+            ),
         ],
     )
     self.assertEqual(
         lm.sample(prompts=['foo', 'bar'], top_k=2, temperature=0.7),
         [
             lm_lib.LMSamplingResult(
-                [lm_lib.LMSample('foo' * 2, score=0.7)]
+                [lm_lib.LMSample('foo' * 2, score=0.7)],
+                usage=lm_lib.LMSamplingUsage(100, 100, 200),
             ),
             lm_lib.LMSamplingResult(
-                [lm_lib.LMSample('bar' * 2, score=0.7)]
+                [lm_lib.LMSample('bar' * 2, score=0.7)],
+                usage=lm_lib.LMSamplingUsage(100, 100, 200),
             ),
         ],
     )
@@ -144,6 +170,8 @@ class LanguageModelTest(unittest.TestCase):
     response = lm(prompt='foo')
     self.assertEqual(response.text, 'foo')
     self.assertEqual(response.score, -1.0)
+    self.assertIsNone(response.logprobs)
+    self.assertEqual(response.usage, lm_lib.LMSamplingUsage(100, 100, 200))
 
     # Test override sampling_options.
     self.assertEqual(
@@ -158,11 +186,24 @@ class LanguageModelTest(unittest.TestCase):
     self.assertEqual(
         lm.sample(prompts=['foo', 'bar']),
         [
-            lm_lib.LMSamplingResult([lm_lib.LMSample(
-                message_lib.AIMessage('foo', cache_seed=0), score=-1.0)]),
-            lm_lib.LMSamplingResult([lm_lib.LMSample(
-                message_lib.AIMessage('bar', cache_seed=0), score=-1.0)]),
-        ])
+            lm_lib.LMSamplingResult(
+                [
+                    lm_lib.LMSample(
+                        message_lib.AIMessage('foo', cache_seed=0), score=-1.0
+                    )
+                ],
+                usage=lm_lib.LMSamplingUsage(100, 100, 200),
+            ),
+            lm_lib.LMSamplingResult(
+                [
+                    lm_lib.LMSample(
+                        message_lib.AIMessage('bar', cache_seed=0), score=-1.0
+                    )
+                ],
+                usage=lm_lib.LMSamplingUsage(100, 100, 200),
+            ),
+        ],
+    )
     self.assertEqual(cache.stats.num_queries, 2)
     self.assertEqual(cache.stats.num_hits, 0)
     self.assertEqual(cache.stats.num_updates, 2)
@@ -181,10 +222,22 @@ class LanguageModelTest(unittest.TestCase):
     self.assertEqual(
         lm.sample(prompts=['foo', 'baz'], temperature=1.0),
         [
-            lm_lib.LMSamplingResult([lm_lib.LMSample(
-                message_lib.AIMessage('foo', cache_seed=0), score=1.0)]),
-            lm_lib.LMSamplingResult([lm_lib.LMSample(
-                message_lib.AIMessage('baz', cache_seed=0), score=1.0)]),
+            lm_lib.LMSamplingResult(
+                [
+                    lm_lib.LMSample(
+                        message_lib.AIMessage('foo', cache_seed=0), score=1.0
+                    )
+                ],
+                usage=lm_lib.LMSamplingUsage(100, 100, 200),
+            ),
+            lm_lib.LMSamplingResult(
+                [
+                    lm_lib.LMSample(
+                        message_lib.AIMessage('baz', cache_seed=0), score=1.0
+                    )
+                ],
+                usage=lm_lib.LMSamplingUsage(100, 100, 200),
+            ),
         ],
     )
     self.assertEqual(cache.stats.num_queries, 6)

langfun/core/llms/cache/in_memory_test.py

@@ -62,10 +62,19 @@ class InMemoryLMCacheTest(unittest.TestCase):
 
     def cache_entry(response_text, cache_seed=0):
       return base.LMCacheEntry(
-          lf.LMSamplingResult([
-              lf.LMSample(
-                  lf.AIMessage(response_text, cache_seed=cache_seed), score=1.0)
-          ])
+          lf.LMSamplingResult(
+              [
+                  lf.LMSample(
+                      lf.AIMessage(response_text, cache_seed=cache_seed),
+                      score=1.0
+                  )
+              ],
+              usage=lf.LMSamplingUsage(
+                  1,
+                  len(response_text),
+                  len(response_text) + 1,
+              )
+          )
       )
 
     self.assertEqual(

langfun/core/llms/fake.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Fake LMs for testing."""
 
+import abc
 from typing import Annotated
 import langfun.core as lf
 
@@ -23,15 +24,32 @@ class Fake(lf.LanguageModel):
   def _score(self, prompt: lf.Message, completions: list[lf.Message]):
     return [lf.LMScoringResult(score=-i * 1.0) for i in range(len(completions))]
 
+  def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
+    results = []
+    for prompt in prompts:
+      response = self._response_from(prompt)
+      results.append(
+          lf.LMSamplingResult(
+              [lf.LMSample(response, 1.0)],
+              usage=lf.LMSamplingUsage(
+                  prompt_tokens=len(prompt.text),
+                  completion_tokens=len(response.text),
+                  total_tokens=len(prompt.text) + len(response.text),
+              )
+          )
+      )
+    return results
+
+  @abc.abstractmethod
+  def _response_from(self, prompt: lf.Message) -> lf.Message:
+    """Returns the response for the given prompt."""
+
 
 class Echo(Fake):
   """A simple echo language model for testing."""
 
-  def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
-    return [
-        lf.LMSamplingResult([lf.LMSample(prompt.text, 1.0)])
-        for prompt in prompts
-    ]
+  def _response_from(self, prompt: lf.Message) -> lf.Message:
+    return lf.AIMessage(prompt.text)
 
 
 @lf.use_init_args(['response'])
@@ -43,11 +61,8 @@ class StaticResponse(Fake):
       'A canned response that will be returned regardless of the prompt.'
   ]
 
-  def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
-    return [
-        lf.LMSamplingResult([lf.LMSample(self.response, 1.0)])
-        for _ in prompts
-    ]
+  def _response_from(self, prompt: lf.Message) -> lf.Message:
+    return lf.AIMessage(self.response)
 
 
 @lf.use_init_args(['mapping'])
@@ -59,11 +74,8 @@ class StaticMapping(Fake):
       'A mapping from prompt to response.'
   ]
 
-  def _sample(self, prompts: list[str]) -> list[lf.LMSamplingResult]:
-    return [
-        lf.LMSamplingResult([lf.LMSample(self.mapping[prompt], 1.0)])
-        for prompt in prompts
-    ]
+  def _response_from(self, prompt: lf.Message) -> lf.Message:
+    return lf.AIMessage(self.mapping[prompt])
 
 
 @lf.use_init_args(['sequence'])
@@ -79,10 +91,7 @@ class StaticSequence(Fake):
     super()._on_bound()
     self._pos = 0
 
-  def _sample(self, prompts: list[str]) -> list[lf.LMSamplingResult]:
-    results = []
-    for _ in prompts:
-      results.append(lf.LMSamplingResult(
-          [lf.LMSample(self.sequence[self._pos], 1.0)]))
-      self._pos += 1
-    return results
+  def _response_from(self, prompt: lf.Message) -> lf.Message:
+    r = lf.AIMessage(self.sequence[self._pos])
+    self._pos += 1
+    return r

langfun/core/llms/fake_test.py

@@ -25,7 +25,12 @@ class EchoTest(unittest.TestCase):
   def test_sample(self):
     lm = fakelm.Echo()
     self.assertEqual(
-        lm.sample(['hi']), [lf.LMSamplingResult([lf.LMSample('hi', 1.0)])]
+        lm.sample(['hi']),
+        [
+            lf.LMSamplingResult(
+                [lf.LMSample('hi', 1.0)],
+                lf.LMSamplingUsage(2, 2, 4))
+        ]
     )
 
   def test_call(self):
@@ -53,11 +58,21 @@ class StaticResponseTest(unittest.TestCase):
     lm = fakelm.StaticResponse(canned_response)
     self.assertEqual(
         lm.sample(['hi']),
-        [lf.LMSamplingResult([lf.LMSample(canned_response, 1.0)])],
+        [
+            lf.LMSamplingResult(
+                [lf.LMSample(canned_response, 1.0)],
+                usage=lf.LMSamplingUsage(2, 38, 40)
+            )
+        ],
     )
     self.assertEqual(
         lm.sample(['Tell me a joke.']),
-        [lf.LMSamplingResult([lf.LMSample(canned_response, 1.0)])],
+        [
+            lf.LMSamplingResult(
+                [lf.LMSample(canned_response, 1.0)],
+                usage=lf.LMSamplingUsage(15, 38, 53)
+            )
+        ],
     )
 
   def test_call(self):
@@ -85,8 +100,14 @@ class StaticMappingTest(unittest.TestCase):
     self.assertEqual(
         lm.sample(['Hi', 'How are you?']),
         [
-            lf.LMSamplingResult([lf.LMSample('Hello', 1.0)]),
-            lf.LMSamplingResult([lf.LMSample('I am fine, how about you?', 1.0)])
+            lf.LMSamplingResult(
+                [lf.LMSample('Hello', 1.0)],
+                usage=lf.LMSamplingUsage(2, 5, 7)
+            ),
+            lf.LMSamplingResult(
+                [lf.LMSample('I am fine, how about you?', 1.0)],
+                usage=lf.LMSamplingUsage(12, 25, 37)
+            )
         ]
     )
     with self.assertRaises(KeyError):
@@ -104,8 +125,14 @@ class StaticSequenceTest(unittest.TestCase):
     self.assertEqual(
         lm.sample(['Hi', 'How are you?']),
         [
-            lf.LMSamplingResult([lf.LMSample('Hello', 1.0)]),
-            lf.LMSamplingResult([lf.LMSample('I am fine, how about you?', 1.0)])
+            lf.LMSamplingResult(
+                [lf.LMSample('Hello', 1.0)],
+                usage=lf.LMSamplingUsage(2, 5, 7)
+            ),
+            lf.LMSamplingResult(
+                [lf.LMSample('I am fine, how about you?', 1.0)],
+                usage=lf.LMSamplingUsage(12, 25, 37)
+            )
         ]
     )
     with self.assertRaises(IndexError):

langfun/core/llms/openai.py

@@ -26,20 +26,6 @@ from openai import openai_object
 import pyglove as pg
 
 
-class Usage(pg.Object):
-  """Usage information per completion."""
-
-  prompt_tokens: int
-  completion_tokens: int
-  total_tokens: int
-
-
-class LMSamplingResult(lf.LMSamplingResult):
-  """LMSamplingResult with usage information."""
-
-  usage: Usage | None = None
-
-
 SUPPORTED_MODELS_AND_SETTINGS = [
     # Model name, max concurrent requests.
     # The concurrent requests is estimated by TPM/RPM from
@@ -181,7 +167,7 @@ class OpenAI(lf.LanguageModel):
       args['stop'] = options.stop
     return args
 
-  def _sample(self, prompts: list[lf.Message]) -> list[LMSamplingResult]:
+  def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
     assert self._api_initialized
     if self.is_chat_model:
       return self._chat_complete_batch(prompts)
@@ -189,7 +175,8 @@ class OpenAI(lf.LanguageModel):
       return self._complete_batch(prompts)
 
   def _complete_batch(
-      self, prompts: list[lf.Message]) -> list[LMSamplingResult]:
+      self, prompts: list[lf.Message]
+  ) -> list[lf.LMSamplingResult]:
 
     def _open_ai_completion(prompts):
       response = openai.Completion.create(
@@ -204,13 +191,13 @@ class OpenAI(lf.LanguageModel):
             lf.LMSample(choice.text.strip(), score=choice.logprobs or 0.0)
         )
 
-      usage = Usage(
+      usage = lf.LMSamplingUsage(
           prompt_tokens=response.usage.prompt_tokens,
           completion_tokens=response.usage.completion_tokens,
           total_tokens=response.usage.total_tokens,
       )
       return [
-          LMSamplingResult(
+          lf.LMSamplingResult(
               samples_by_index[index], usage=usage if index == 0 else None
           )
           for index in sorted(samples_by_index.keys())
@@ -231,7 +218,7 @@ class OpenAI(lf.LanguageModel):
 
   def _chat_complete_batch(
       self, prompts: list[lf.Message]
-  ) -> list[LMSamplingResult]:
+  ) -> list[lf.LMSamplingResult]:
     def _open_ai_chat_completion(prompt: lf.Message):
       if self.multimodal:
         content = []
@@ -272,9 +259,9 @@ class OpenAI(lf.LanguageModel):
             )
         )
 
-      return LMSamplingResult(
+      return lf.LMSamplingResult(
           samples=samples,
-          usage=Usage(
+          usage=lf.LMSamplingUsage(
               prompt_tokens=response.usage.prompt_tokens,
               completion_tokens=response.usage.completion_tokens,
               total_tokens=response.usage.total_tokens,

langfun/core/llms/openai_test.py

@@ -32,11 +32,14 @@ def mock_completion_query(prompt, *, n=1, **kwargs):
           text=f'Sample {k} for prompt {i}.',
           logprobs=k / 10,
       ))
-  return pg.Dict(choices=choices, usage=openai.Usage(
-      prompt_tokens=100,
-      completion_tokens=100,
-      total_tokens=200,
-  ))
+  return pg.Dict(
+      choices=choices,
+      usage=lf.LMSamplingUsage(
+          prompt_tokens=100,
+          completion_tokens=100,
+          total_tokens=200,
+      ),
+  )
 
 
 def mock_chat_completion_query(messages, *, n=1, **kwargs):
@@ -49,11 +52,14 @@ def mock_chat_completion_query(messages, *, n=1, **kwargs):
         ),
         logprobs=None,
     ))
-  return pg.Dict(choices=choices, usage=openai.Usage(
-      prompt_tokens=100,
-      completion_tokens=100,
-      total_tokens=200,
-  ))
+  return pg.Dict(
+      choices=choices,
+      usage=lf.LMSamplingUsage(
+          prompt_tokens=100,
+          completion_tokens=100,
+          total_tokens=200,
+      ),
+  )
 
 
 def mock_chat_completion_query_vision(messages, *, n=1, **kwargs):
@@ -69,11 +75,14 @@ def mock_chat_completion_query_vision(messages, *, n=1, **kwargs):
         ),
         logprobs=None,
     ))
-  return pg.Dict(choices=choices, usage=openai.Usage(
-      prompt_tokens=100,
-      completion_tokens=100,
-      total_tokens=200,
-  ))
+  return pg.Dict(
+      choices=choices,
+      usage=lf.LMSamplingUsage(
+          prompt_tokens=100,
+          completion_tokens=100,
+          total_tokens=200,
+      ),
+  )
 
 
 class OpenaiTest(unittest.TestCase):
@@ -169,18 +178,28 @@ class OpenaiTest(unittest.TestCase):
       )
 
     self.assertEqual(len(results), 2)
-    self.assertEqual(results[0], openai.LMSamplingResult([
-        lf.LMSample('Sample 0 for prompt 0.', score=0.0),
-        lf.LMSample('Sample 1 for prompt 0.', score=0.1),
-        lf.LMSample('Sample 2 for prompt 0.', score=0.2),
-    ], usage=openai.Usage(
-        prompt_tokens=100, completion_tokens=100, total_tokens=200)))
-
-    self.assertEqual(results[1], openai.LMSamplingResult([
-        lf.LMSample('Sample 0 for prompt 1.', score=0.0),
-        lf.LMSample('Sample 1 for prompt 1.', score=0.1),
-        lf.LMSample('Sample 2 for prompt 1.', score=0.2),
-    ]))
+    self.assertEqual(
+        results[0],
+        lf.LMSamplingResult(
+            [
+                lf.LMSample('Sample 0 for prompt 0.', score=0.0),
+                lf.LMSample('Sample 1 for prompt 0.', score=0.1),
+                lf.LMSample('Sample 2 for prompt 0.', score=0.2),
+            ],
+            usage=lf.LMSamplingUsage(
+                prompt_tokens=100, completion_tokens=100, total_tokens=200
+            ),
+        ),
+    )
+
+    self.assertEqual(
+        results[1],
+        lf.LMSamplingResult([
+            lf.LMSample('Sample 0 for prompt 1.', score=0.0),
+            lf.LMSample('Sample 1 for prompt 1.', score=0.1),
+            lf.LMSample('Sample 2 for prompt 1.', score=0.2),
+        ]),
+    )
 
   def test_sample_chat_completion(self):
     with mock.patch('openai.ChatCompletion.create') as mock_chat_completion:
@@ -191,18 +210,32 @@ class OpenaiTest(unittest.TestCase):
       )
 
     self.assertEqual(len(results), 2)
-    self.assertEqual(results[0], openai.LMSamplingResult([
-        lf.LMSample('Sample 0 for message.', score=0.0),
-        lf.LMSample('Sample 1 for message.', score=0.0),
-        lf.LMSample('Sample 2 for message.', score=0.0),
-    ], usage=openai.Usage(
-        prompt_tokens=100, completion_tokens=100, total_tokens=200)))
-    self.assertEqual(results[1], openai.LMSamplingResult([
-        lf.LMSample('Sample 0 for message.', score=0.0),
-        lf.LMSample('Sample 1 for message.', score=0.0),
-        lf.LMSample('Sample 2 for message.', score=0.0),
-    ], usage=openai.Usage(
-        prompt_tokens=100, completion_tokens=100, total_tokens=200)))
+    self.assertEqual(
+        results[0],
+        lf.LMSamplingResult(
+            [
+                lf.LMSample('Sample 0 for message.', score=0.0),
+                lf.LMSample('Sample 1 for message.', score=0.0),
+                lf.LMSample('Sample 2 for message.', score=0.0),
+            ],
+            usage=lf.LMSamplingUsage(
+                prompt_tokens=100, completion_tokens=100, total_tokens=200
+            ),
+        ),
+    )
+    self.assertEqual(
+        results[1],
+        lf.LMSamplingResult(
+            [
+                lf.LMSample('Sample 0 for message.', score=0.0),
+                lf.LMSample('Sample 1 for message.', score=0.0),
+                lf.LMSample('Sample 2 for message.', score=0.0),
+            ],
+            usage=lf.LMSamplingUsage(
+                prompt_tokens=100, completion_tokens=100, total_tokens=200
+            ),
+        ),
+    )
 
   def test_sample_with_contextual_options(self):
     with mock.patch('openai.Completion.create') as mock_completion:
@@ -212,11 +245,18 @@ class OpenaiTest(unittest.TestCase):
         results = lm.sample(['hello'])
 
     self.assertEqual(len(results), 1)
-    self.assertEqual(results[0], openai.LMSamplingResult([
-        lf.LMSample('Sample 0 for prompt 0.', score=0.0),
-        lf.LMSample('Sample 1 for prompt 0.', score=0.1),
-    ], usage=openai.Usage(
-        prompt_tokens=100, completion_tokens=100, total_tokens=200)))
+    self.assertEqual(
+        results[0],
+        lf.LMSamplingResult(
+            [
+                lf.LMSample('Sample 0 for prompt 0.', score=0.0),
+                lf.LMSample('Sample 1 for prompt 0.', score=0.1),
+            ],
+            usage=lf.LMSamplingUsage(
+                prompt_tokens=100, completion_tokens=100, total_tokens=200
+            ),
+        ),
+    )
 
 
 if __name__ == '__main__':