| 257 | |
| 258 | |
| 259 | class GPT_Insert(LLM): |
| 260 | |
| 261 | def __init__(self, config, needs_confirmation=False, disable_tqdm=True): |
| 262 | """Initializes the model.""" |
| 263 | self.config = config |
| 264 | self.needs_confirmation = needs_confirmation |
| 265 | self.disable_tqdm = disable_tqdm |
| 266 | |
| 267 | def confirm_cost(self, texts, n, max_tokens): |
| 268 | total_estimated_cost = 0 |
| 269 | for text in texts: |
| 270 | total_estimated_cost += gpt_get_estimated_cost( |
| 271 | self.config, text, max_tokens) * n |
| 272 | print(f"Estimated cost: ${total_estimated_cost:.2f}") |
| 273 | # Ask the user to confirm in the command line |
| 274 | if os.getenv("LLM_SKIP_CONFIRM") is None: |
| 275 | confirm = input("Continue? (y/n) ") |
| 276 | if confirm != 'y': |
| 277 | raise Exception("Aborted.") |
| 278 | |
| 279 | def auto_reduce_n(self, fn, prompt, n): |
| 280 | """Reduces n by half until the function succeeds.""" |
| 281 | try: |
| 282 | return fn(prompt, n) |
| 283 | except BatchSizeException as e: |
| 284 | if n == 1: |
| 285 | raise e |
| 286 | return self.auto_reduce_n(fn, prompt, n // 2) + self.auto_reduce_n(fn, prompt, n // 2) |
| 287 | |
| 288 | def generate_text(self, prompt, n): |
| 289 | if not isinstance(prompt, list): |
| 290 | prompt = [prompt] |
| 291 | if self.needs_confirmation: |
| 292 | self.confirm_cost( |
| 293 | prompt, n, self.config['gpt_config']['max_tokens']) |
| 294 | batch_size = self.config['batch_size'] |
| 295 | assert batch_size == 1 |
| 296 | prompt_batches = [prompt[i:i + batch_size] |
| 297 | for i in range(0, len(prompt), batch_size)] |
| 298 | if not self.disable_tqdm: |
| 299 | print( |
| 300 | f"[{self.config['name']}] Generating {len(prompt) * n} completions, split into {len(prompt_batches)} batches of (maximum) size {batch_size * n}") |
| 301 | text = [] |
| 302 | for prompt_batch in tqdm(prompt_batches, disable=self.disable_tqdm): |
| 303 | text += self.auto_reduce_n(self.__generate_text, prompt_batch, n) |
| 304 | return text |
| 305 | |
| 306 | def log_probs(self, text, log_prob_range=None): |
| 307 | raise NotImplementedError |
| 308 | |
| 309 | def __generate_text(self, prompt, n): |
| 310 | """Generates text from the model.""" |
| 311 | config = self.config['gpt_config'].copy() |
| 312 | config['n'] = n |
| 313 | # Split prompts into prefixes and suffixes with the [APE] token (do not include the [APE] token in the suffix) |
| 314 | prefix = prompt[0].split('[APE]')[0] |
| 315 | suffix = prompt[0].split('[APE]')[1] |
| 316 | response = None |