[evals] Refactor evals package to expose completion_fn.

PAIR=jasonwei

Co-authored-by: Jason Wei <[email protected]>

evals/__init__.py

@@ -1,4 +1,4 @@
-from .api import check_sampled_text, completion_query, sample_freeform
+from .api import check_sampled_text, completion_query, sample_freeform, postprocess_sample_freeform, record_and_check_match
 from .base import ModelSpec, ModelSpecs
 from .data import get_csv, get_json, get_jsonl, get_jsonls, get_lines, iter_jsonls
 from .eval import Eval

evals/api.py

@@ -97,6 +97,7 @@ def completion_query(
  return result, openai_create_prompt, metadata
 
 
+# TODO(hwc): remove this
 def check_sampled_text(
  model_spec: ModelSpec,
  prompt: Union[OpenAICreatePrompt, OpenAICreateChatPrompt, Prompt],
@@ -123,13 +124,6 @@ def check_sampled_text(
  =======
  The option that was picked, i.e., matched the completion, or None.
  """
- if isinstance(expected, tuple):
- expected = list(expected)
- elif not isinstance(expected, list):
- expected = [expected]
- if options is None:
- options = expected
-
  result, actual_prompt, metadata = completion_query(
  prompt=prompt,
  temperature=0.0,
@@ -139,6 +133,31 @@ def check_sampled_text(
 
  sampled = choice["text"].strip() if model_spec.strip_completion else choice["text"]
 
+ return record_and_check_match(
+ prompt=actual_prompt,
+ sampled=sampled,
+ expected=expected,
+ metadata=metadata,
+ separator=separator,
+ options=options,
+ )
+
+
+def record_and_check_match(
+ prompt: Union[OpenAICreatePrompt, OpenAICreateChatPrompt],
+ sampled: str,
+ expected: Union[str, list[str], tuple[str]],
+ metadata: dict,
+ separator: Callable[[str], bool] = None,
+ options: Optional[list[str]] = None,
+):
+ if isinstance(expected, tuple):
+ expected = list(expected)
+ elif not isinstance(expected, list):
+ expected = [expected]
+ if options is None:
+ options = expected
+
  picked = None
  for option in options:
  if not sampled.startswith(option):
@@ -153,7 +172,7 @@ def check_sampled_text(
  break
 
  result = {
- "prompt": actual_prompt,
+ "prompt": prompt,
  "sampled": sampled,
  "options": options,
  "picked": picked,
@@ -175,7 +194,7 @@ def sample_freeform(
  top_p: float = 0.9,
  max_tokens: int = 512,
  stop: Optional[str] = None,
- n_samples: int = None,
+ n_samples: Optional[int] = None,
  return_logprobs: bool = False,
  **kwargs,
 ) -> Union[str, list[str], dict]:
@@ -215,10 +234,51 @@ def sample_freeform(
  headers={},
  **kwargs,
  )
+ return postprocess_sample_freeform(
+ response,
+ actual_prompt,
+ metadata,
+ model_spec,
+ n_samples=n_samples,
+ return_logprobs=return_logprobs,
+ **kwargs)
+
+
+def postprocess_sample_freeform(
+ response: dict,
+ prompt: Union[OpenAICreatePrompt, OpenAICreateChatPrompt, Prompt],
+ metadata: dict,
+ model_spec: ModelSpec,
+ *,
+ n_samples: Optional[int] = None,
+ return_logprobs: bool = False,
+ **kwargs,
+) -> Union[str, list[str], dict]:
+ """
+ Records the sampled response, prompt and metedata, and returns the sampled text.
+ Typically called after `sample_freeform`.
+
+ ARGS
+ ====
+ `response`: The result of the API call.
+ `prompt`: See `completion_query`.
+ `n_samples`: The number of samples to generate (1 if None).
+ `return_logprobs`: If True, returns the tokens and corresponding logprobs
+ in addition to the sampled text.
+ `kwargs`: See `completion_query`.
+
+ RETURNS
+ =======
+ If `return_logprobs` is True, returns a dict with the sampled text, tokens,
+ and corresponding logprobs. If `n_samples` is None, the outer list is
+ removed from all values.
+ Otherwise, returns the sampled text, or a list of sampled texts if
+ `n_samples` is not None.
+ """
  sampled = [choice["text"] for choice in response["choices"]]
  if n_samples is None:
  sampled = sampled[0]
- record_sampling(prompt=actual_prompt, sampled=sampled, metadata=metadata)
+ record_sampling(prompt=prompt, sampled=sampled, metadata=metadata)
 
  if return_logprobs:
  assert not model_spec.is_chat, "logprobs only works for non-chat models"

evals/elsuite/basic/fuzzy_match.py

@@ -11,20 +11,26 @@ def __init__(
  samples_jsonl: str,
  *args,
  max_tokens: int = 500,
+ completion_fn: utils.CompletionFn = evals.completion_query,
  **kwargs,
  ):
  super().__init__(model_specs, *args, **kwargs)
  self.max_tokens = max_tokens
  self.samples_jsonl = samples_jsonl
+ self._completion_fn = completion_fn
 
  def eval_sample(self, test_sample, rng):
+ del rng
  prompt, correct_answers = test_sample["input"], test_sample["ideal"]
- generated_answer = evals.sample_freeform(
- self.model_spec,
- prompt,
- temperature=0.0,
+ response, actual_prompt, metadata = self._completion_fn(
+ prompt=prompt,
+ temperature=0.0, # Q: why are these hardcoded?
  max_tokens=16,
+ model_spec=self.model_spec,
  )
+ generated_answer: str = evals.postprocess_sample_freeform(
+ response, actual_prompt, metadata, self.model_spec)
+
  matches = [
  utils.fuzzy_match(generated_answer, correct_answer)
  for correct_answer in correct_answers
@@ -40,7 +46,7 @@ def eval_sample(self, test_sample, rng):
  )
 
  def run(self, recorder: RecorderBase):
- samples = evals.get_jsonl(self.samples_jsonl)
+ samples = self.get_samples()
  self.eval_all_samples(recorder, samples)
 
  return {

evals/elsuite/basic/includes.py

@@ -1,7 +1,7 @@
 from typing import Any
 
 import evals
-import evals.elsuite.utils
+from evals.elsuite import utils
 import evals.metrics
 import numpy as np
 
@@ -13,24 +13,31 @@ def __init__(
  samples_jsonl: str,
  *args,
  max_tokens: int = 500,
+ completion_fn: utils.CompletionFn = evals.completion_query,
  **kwargs,
  ):
  super().__init__(model_specs, *args, **kwargs)
  self.max_tokens = max_tokens
  self.samples_jsonl = samples_jsonl
+ self._completion_fn = completion_fn
 
  def eval_sample(self, sample: Any, *_):
- sampled = evals.sample_freeform(
- self.model_spec, sample["input"], max_tokens=self.max_tokens
+ response, actual_prompt, metadata = self._completion_fn(
+ prompt=sample["input"],
+ max_tokens=self.max_tokens,
+ model_spec=self.model_spec,
  )
+ sampled: str = evals.postprocess_sample_freeform(
+ response, actual_prompt, metadata, self.model_spec)
+
  includes_answer = any(
- [evals.elsuite.utils.get_answer(sampled, ref) for ref in sample["ideal"]]
+ [utils.get_answer(sampled, ref) for ref in sample["ideal"]]
  )
  evals.record.record_metrics(accuracy=float(includes_answer))
  return includes_answer
 
  def run(self, recorder):
- samples = evals.get_jsonl(self.samples_jsonl)
+ samples = self.get_samples()
  self.eval_all_samples(recorder, samples)
  events = recorder.get_scores("accuracy")
  return {

evals/elsuite/basic/match.py

@@ -2,6 +2,7 @@
 
 import evals
 import evals.metrics
+from evals.elsuite import utils
 from evals.prompt.base import is_chat_prompt
 
 
@@ -14,6 +15,7 @@ def __init__(
  max_tokens: int = 500,
  num_few_shot: int = 0,
  few_shot_jsonl: str = None,
+ completion_fn: utils.CompletionFn = evals.completion_query,
  **kwargs,
  ):
  super().__init__(model_specs, *args, **kwargs)
@@ -24,6 +26,7 @@ def __init__(
  assert few_shot_jsonl is not None, "few shot requires few shot sample dataset"
  self.few_shot_jsonl = few_shot_jsonl
  self.few_shot = evals.get_jsonl(self.few_shot_jsonl)
+ self._completion_fn = completion_fn
 
  def eval_sample(self, sample: Any, *_):
  prompt = sample["input"]
@@ -34,10 +37,23 @@ def eval_sample(self, sample: Any, *_):
  prompt += s["sample"]
  prompt += sample["input"][-1:]
 
- return evals.check_sampled_text(self.model_spec, prompt, expected=sample["ideal"])
+ # TODO(hwc): is there a case where we want to use `result` other than "choices"?
+ result, actual_prompt, metadata = self._completion_fn(
+ prompt=prompt,
+ temperature=0.0,
+ model_spec=self.model_spec,
+ )
+ choice = result["choices"][0]
+ sampled = choice["text"].strip() if self.model_spec.strip_completion else choice["text"]
+ return evals.record_and_check_match(
+ prompt=actual_prompt,
+ sampled=sampled,
+ expected=sample["ideal"],
+ metadata=metadata
+ )
 
  def run(self, recorder):
- samples = evals.get_jsonl(self.samples_jsonl)
+ samples= self.get_samples()
  self.eval_all_samples(recorder, samples)
  events = recorder.get_events("match")
  return {

evals/elsuite/modelgraded/classify.py

@@ -319,7 +319,7 @@ def eval_sample(self, test_sample: dict, rng: Random) -> None:
  return choice
 
  def run(self, recorder):
- samples = evals.get_jsonl(self.samples_jsonl)
+ samples = self.get_samples()
 
  self.eval_all_samples(recorder, samples)
  all_sample_metrics = recorder.get_metrics()

evals/elsuite/translate.py

@@ -4,6 +4,7 @@
 
 import evals
 import evals.metrics
+from evals.elsuite import utils
 from evals.prompt.base import is_chat_prompt
 
 
@@ -16,6 +17,7 @@ def __init__(
  max_tokens: int = 500,
  num_few_shot: int = 0,
  few_shot_jsonl: str = None,
+ completion_fn: utils.CompletionFn = evals.completion_query,
  **kwargs,
  ):
  super().__init__(model_specs, *args, **kwargs)
@@ -29,6 +31,7 @@ def __init__(
  self.few_shot = evals.get_jsonl(self.few_shot_jsonl)
 
  self.bleu = BLEU(effective_order=True)
+ self._completion_fn = completion_fn
 
  def eval_sample(self, sample: Any, *_):
  prompt = sample["input"]
@@ -45,7 +48,13 @@ def eval_sample(self, sample: Any, *_):
  elif not isinstance(expected, list):
  expected = [expected]
 
- sampled = evals.sample_freeform(self.model_spec, prompt, max_tokens=self.max_tokens)
+ response, actual_prompt, metadata = self._completion_fn(
+ prompt=prompt,
+ max_tokens=self.max_tokens,
+ model_spec=self.model_spec,
+ )
+ sampled: str = evals.postprocess_sample_freeform(
+ response, actual_prompt, metadata, self.model_spec)
 
  score = None
  if expected is not None:
@@ -61,7 +70,7 @@ def eval_sample(self, sample: Any, *_):
  return match
 
  def run(self, recorder):
- samples = evals.get_jsonl(self.samples_jsonl)
+ samples = self.get_samples()
  self.eval_all_samples(recorder, samples)
  events = recorder.get_events("match")
 

evals/elsuite/utils.py

@@ -2,10 +2,19 @@
 import re
 import string
 from collections import Counter, defaultdict
+from typing import Union
+from typing_extensions import Protocol
 
 from evals.api import sample_freeform
 from evals.prompt.base import chat_prompt_to_text_prompt, is_chat_prompt
 
+from evals.base import ModelSpec
+from evals.prompt.base import (
+ OpenAICreateChatPrompt,
+ OpenAICreatePrompt,
+ Prompt,
+)
+
 
 def get_answer(text, answer_prompt):
  idx = text.rfind(answer_prompt)
@@ -135,3 +144,31 @@ def __call__(self, **kwargs):
  **self.completion_kwargs,
  )
  return completion, prompt
+
+
+class CompletionFn(Protocol):
+
+ def __call__(
+ self,
+ model_spec: ModelSpec,
+ prompt: Union[OpenAICreatePrompt, OpenAICreateChatPrompt, Prompt],
+ **kwargs
+ ) ->tuple[dict, Union[OpenAICreatePrompt, OpenAICreateChatPrompt], dict]:
+ """
+ ARGS
+ ====
+ `model_spec`: `ModelSpec` containing model details to use in the query.
+ This should be the dict returned by `registry.get_model()`.
+ If `model_spec` is not provided, we use the default model that was
+ intialized at the beginning of the run.
+ `prompt`: Either a `Prompt` object or a raw prompt that will get wrapped in
+ the approriate `Prompt` class.
+ `kwargs`: Other arguments passed to the API.
+
+ RETURNS
+ =======
+ The result of the API call.
+ The prompt that was fed into the API call as a str.
+ A dict containing metadata about the query.
+ """
+ pass