fix: dummy client to output tokens and random responses

manifest/clients/dummy.py

@@ -1,6 +1,10 @@
 """Dummy client."""
+import hashlib
 import logging
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+import tiktoken
 
 from manifest.clients.client import Client
 from manifest.request import LMChatRequest, LMRequest, LMScoreRequest, Request
@@ -14,7 +18,13 @@ class DummyClient(Client):
 
  # User param -> (client param, default value)
  PARAMS = {
- "n": ("num_results", 1),
+ "engine": ("model", "text-davinci-003"),
+ "temperature": ("temperature", 0.0),
+ "max_tokens": ("max_tokens", 10),
+ "n": ("n", 1),
+ "top_p": ("top_p", 1.0),
+ "top_k": ("best_of", 1),
+ "batch_size": ("batch_size", 20),
  }
  REQUEST_CLS = LMRequest
  NAME = "dummy"
@@ -33,6 +43,9 @@ def connect(
  connection_str: connection string.
  client_args: client arguments.
  """
+ # We tiktoken as it is faster than HF for tokenizing
+ # Use any model to create the tokenizer
+ self.encoder = tiktoken.get_encoding("cl100k_base")
  for key in self.PARAMS:
  setattr(self, key, client_args.pop(key, self.PARAMS[key][1]))
 
@@ -74,7 +87,65 @@ def get_model_params(self) -> Dict:
  Returns:
  model params.
  """
- return {"engine": "dummy"}
+ return {"engine": "dummy", "model": getattr(self, "engine")}
+
+ def get_mock_output(
+ self, output_toks: int, is_completion: bool, seed: Optional[int] = None
+ ) -> LMModelChoice:
+ """Return mock model output by generating random tokens."""
+ np.random.seed(seed)
+ random_tokens = np.random.randint(
+ 0, self.encoder.max_token_value + 1, output_toks
+ )
+ response = self.encoder.decode(random_tokens) # type: ignore
+ if is_completion:
+ np.random.seed(seed)
+ random_logprobs = np.random.uniform(
+ low=-2, high=-0.00001, size=output_toks
+ ).tolist()
+ else:
+ # Return all Nones to mimic chat models
+ # OpenAI chat models do not return logprobs
+ random_logprobs = [None] * output_toks
+ return LMModelChoice(
+ text=response,
+ token_logprobs=random_logprobs,
+ tokens=random_tokens.tolist(),
+ )
+
+ def get_mock_choices(
+ self,
+ prompt_list: List[str],
+ request_params: Dict,
+ is_completion: bool,
+ ) -> Tuple[List[LMModelChoice], List[Usage]]:
+ """Get choices and usages of mock output."""
+ choices = []
+ usages = []
+ for prompt in prompt_list:
+ num_prompt_tokens = len(self.encoder.encode(prompt))
+ if request_params["temperature"] == 0:
+ # Get integer seed from hash of prompt
+ seed = (
+ int(hashlib.sha256(prompt.encode("utf-8")).hexdigest(), 16)
+ % 10**8
+ )
+ else:
+ # Get random seed
+ seed = None
+ for _ in range(int(request_params["n"])):
+ choice = self.get_mock_output(
+ request_params["max_tokens"], is_completion=is_completion, seed=seed
+ )
+ choices.append(choice)
+ usages.append(
+ Usage(
+ prompt_tokens=num_prompt_tokens,
+ completion_tokens=request_params["max_tokens"],
+ total_tokens=num_prompt_tokens + request_params["max_tokens"],
+ )
+ )
+ return choices, usages
 
  def run_request(self, request: Request) -> Response:
  """
@@ -88,32 +159,19 @@ def run_request(self, request: Request) -> Response:
  request parameters as dict.
  """
  if isinstance(request.prompt, list):
- num_results = len(request.prompt)
+ prompt_list = request.prompt
  else:
- num_results = 1
+ prompt_list = [request.prompt]
  request_params = request.to_dict(self.PARAMS)
 
+ choices, usages = self.get_mock_choices(
+ prompt_list, request_params, is_completion=True
+ )
  return Response(
- response=ModelChoices(
- choices=[LMModelChoice(text="hello")] # type: ignore
- * int(request_params["num_results"])
- * num_results
- ),
+ response=ModelChoices(choices=choices), # type: ignore
  cached=False,
  request=request,
- usages=Usages(
- usages=[
- Usage(
- **{
- "prompt_tokens": 1,
- "completion_tokens": 1,
- "total_tokens": 2,
- }
- )
- ]
- * int(request_params["num_results"])
- * num_results
- ),
+ usages=Usages(usages=usages),
  response_type="text",
  request_type=self.REQUEST_CLS,
  )
@@ -145,35 +203,17 @@ def run_chat_request(
  Returns:
  response.
  """
- num_results = 1
- response_dict = {
- "choices": [
- {
- "text": request.prompt[0]["content"],
- }
- for i in range(num_results)
- ]
- }
+ prompt_list = ["_".join(pmp["content"] for pmp in request.prompt)]
+ request_params = request.to_dict(self.PARAMS)
+
+ choices, usages = self.get_mock_choices(
+ prompt_list, request_params, is_completion=False
+ )
  return Response(
- response=ModelChoices(
- choices=[
- LMModelChoice(**choice) # type: ignore
- for choice in response_dict["choices"]
- ]
- ),
+ response=ModelChoices(choices=choices), # type: ignore
  cached=False,
  request=request,
- usages=Usages(
- usages=[
- Usage(
- **{
- "prompt_tokens": 1,
- "completion_tokens": 1,
- "total_tokens": 2,
- }
- )
- ]
- ),
+ usages=Usages(usages=usages),
  response_type="text",
  request_type=LMChatRequest,
  )
@@ -193,30 +233,19 @@ def run_score_prompt_request(
  request parameters as dict.
  """
  if isinstance(request.prompt, list):
- num_results = len(request.prompt)
+ prompt_list = request.prompt
  else:
- num_results = 1
- response_dict = {
- "choices": [
- {
- "text": request.prompt
- if isinstance(request.prompt, str)
- else request.prompt[i],
- "token_logprobs": [0.3],
- }
- for i in range(num_results)
- ]
- }
+ prompt_list = [request.prompt]
+ request_params = request.to_dict(self.PARAMS)
+
+ choices, usages = self.get_mock_choices(
+ prompt_list, request_params, is_completion=True
+ )
  return Response(
- response=ModelChoices(
- choices=[
- LMModelChoice(**choice) # type: ignore
- for choice in response_dict["choices"]
- ]
- ),
+ response=ModelChoices(choices=choices), # type: ignore
  cached=False,
  request=request,
- usages=None,
+ usages=Usages(usages=usages),
  response_type="text",
  request_type=LMScoreRequest,
  )

manifest/response.py

@@ -53,7 +53,7 @@ class LMModelChoice(BaseModel):
  """Model single completion."""
 
  text: str
- token_logprobs: Optional[List[float]] = None
+ token_logprobs: Optional[List[Optional[float]]] = None
  tokens: Optional[List[str]] = None