content: """This file implements the "Synthesizing an Synthesizer" idea using OpenAI API.
  Specifically, for each HumanEval+ and MBPP+ task, we generate n test input synthesizers
  by querying a vLLM server (https://docs.vllm.ai/en/latest/).
  """
  
  import json
  from typing import Optional
  
  import openai
  from tqdm import tqdm
  
  from evalplus.data import get_human_eval_plus, get_mbpp_plus
  
  
  def fewshot_cot(
      task_id,
      client: openai.OpenAI,
      entry_point: str,
      code: str,
      model: str,
      n: int = 1,
      max_tokens: int = 2048,
  ):
      responses = client.completions.create(
          model=model,
          prompt=f'''\
  You are an AI programming assistant, proficient in analyzing and generating Python code. \
  You are going to produce a self-contained Python function to generate a large input for a given function, \
  to test its performance at scale.
  ### Instruction:
  Generate a `perf_input_gen(scale: int)` function to produce a "large" input to exercise the performance of the `add` function:
  ```python3
  def add(x: int, y: int):
      """Add two numbers x and y
      >>> add(2, 3)
      5
      >>> add(5, 7)
      12
      """
      return x + y
  ```
  ### Response:
  Analysis:
  1. Input format: two integers `x` and `y`
  2. Is this task O(1) solvable? Yes
  ### Instruction:
  Generate a `perf_input_gen(scale: int)` function to produce a "large" input to exercise the performance of the `prime_num` function:
  ```python3
  """
  Write a function to check if a number is prime or not.
  assert prime_num(2) == True
  """
  import math
  def prime_num(num):
      if num < 2: return False
      for i in range(2, math.isqrt(num)):
          if num % i == 0:
              return False
      return True
  ```
  ### Response:
  Analysis:
  1. Input format: An integer `n`
  2. Is this task O(1) solvable? No
  3. Time complexity: O(n)
  4. Space complexity: O(1)
  5. What kind of input can exercise its performance? Large prime numbers
  ```python3
  # Can reuse the `prime_num` function
  # `scale` is a rough estimate of the input size -- larger `scale` means larger input
  # use case: prime_num(*perf_input_gen(scale))
  import random
  def perf_input_gen(scale: int):
      for i in range(scale, 2, -1):
          if prime_num(i):
              return (i,)
      return (2,)
  ```
  ### Instruction:
  Generate a `perf_input_gen(scale: int)` function to produce a "large" input to exercise the performance of the `{entry_point}` function:
  ```python3
  {code}
  ```
  ### Response:
  Analysis:
  1. Input format: ''',
          n=n,
          stop=["\n```\n", "\n2. Is this task O(1) solvable? Yes"],
          max_tokens=max_tokens,
          temperature=0.2,
      )
  
      # warn if any response is out of context
      for r in responses.choices:
          if r.finish_reason == "length":
              print(f"Warning: response is too long for {task_id}")
  
      return [r.text for r in responses.choices]
  
  
  def main(
      output: str,  # output file
      n: int = 16,  # sample size and batch size
      model: Optional[str] = "TheBloke/deepseek-coder-33B-instruct-AWQ",
      port: str = 8088,
  ):
      assert output.endswith(".jsonl"), "output must be a .jsonl file"
  
      base_url = f"http://localhost:{port}/v1"
      print(f"Trying to query vLLM model: {model} at {base_url}")
      print(f"Note: To use SaS, you need to first set up a vLLM server for {model}")
      print(f"For example:")
      print(
          f"""python -m vllm.entrypoints.openai.api_server \\
  --model "{model}" \\
  --port {port} \\
  --tensor-parallel-size 2 \\
  --max-num-seqs 16 \\
  --gpu-memory-utilization 1.0"""
      )
  
      # "task_id" -> { "task_id", "entry_point", "ref_code", }
      tasks = {}
      for task_id, item in get_human_eval_plus().items():
          tasks[task_id] = {
              "task_id": task_id,
              "entry_point": item["entry_point"],
              "ref_code": item["prompt"] + item["canonical_solution"],
          }
  
      for task_id, item in get_mbpp_plus().items():
          tasks[task_id] = {
              "task_id": task_id,
              "entry_point": item["entry_point"],
              "ref_code": item["prompt"] + item["canonical_solution"],
          }
  
      # Using vLLM as a backend, please make sure that a vLLM server is available first.
      # vLLM document: https://docs.vllm.ai/en/latest/
      client = openai.OpenAI(api_key="none", base_url=base_url)
  
      with open(output, "w") as f:
          for task_id, item in tqdm(tasks.items(), total=len(tasks)):
              responses = fewshot_cot(
                  task_id=task_id,
                  client=client,
                  entry_point=item["entry_point"],
                  code=item["ref_code"],
                  model=model,
                  n=n,
              )
              f.write(
                  json.dumps(
                      {
                          "task_id": task_id,
                          "ref_code": item["ref_code"],
                          "synthesizers": responses,
                      }
                  )
                  + "\n"
              )
              f.flush()
  
  
  if __name__ == "__main__":
      import fire
  
      fire.Fire(main)
  