Log model and usage stats in sample logger

openai · Jan 3, 2024 · 27a9d2b · 27a9d2b
1 parent ded9382
commit 27a9d2b
Showing 1 changed file with 12 additions and 2 deletions.
diff --git a/evals/completion_fns/openai.py b/evals/completion_fns/openai.py
@@ -88,7 +88,12 @@ def __call__(
             **{**kwargs, **self.extra_options},
         )
         result = OpenAICompletionResult(raw_data=result, prompt=openai_create_prompt)
-        record_sampling(prompt=result.prompt, sampled=result.get_completions())
+        record_sampling(
+            prompt=result.prompt,
+            sampled=result.get_completions(),
+            model=result.raw_data.model,
+            usage=result.raw_data.usage,
+        )
         return result
 
 
@@ -133,5 +138,10 @@ def __call__(
             **{**kwargs, **self.extra_options},
         )
         result = OpenAIChatCompletionResult(raw_data=result, prompt=openai_create_prompt)
-        record_sampling(prompt=result.prompt, sampled=result.get_completions())
+        record_sampling(
+            prompt=result.prompt,
+            sampled=result.get_completions(),
+            model=result.raw_data.model,
+            usage=result.raw_data.usage,
+        )
         return result