import os
import json
import tqdm
import openai
from time import sleep
from pathlib import Path

chat_models = ["gpt-4", "gpt-3.5-turbo"]
message_history = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "Write a unique, surprising, extremely randomized story with highly unpredictable changes of events."}
]

completion_models = ["text-davinci-003", "text-davinci-001", "davinci-instruct-beta", "davinci"]
prompt = "[System: You are a helpful assistant]\n\nUser: Write a unique, surprising, extremely randomized story with highly unpredictable changes of events.\n\nAI:"

results = []

import time
class TimeIt:
    def __init__(self, name): self.name = name
    def __enter__(self): self.start = time.time()
    def __exit__(self, *args): print(f"{self.name} took {time.time() - self.start} seconds")

C = 30  # number of completions to make per model
N = 128 # max_tokens
# Testing chat models
for model in chat_models:
    sequences = set()
    errors = 0 # although I track errors, at no point were any errors ever emitted
    with TimeIt(model):
        for _ in range(C):
            try:
                completion = openai.ChatCompletion.create(
                    model=model,
                    messages=message_history,
                    max_tokens=N,
                    temperature=0,
                    logit_bias={"100257": -100.0}, # this doesn't really do anything, because chat models don't do <|endoftext|> much
                )
                sequences.add(completion.choices[0].message['content'])
                sleep(1) # cheaply avoid rate limiting
            except Exception as e:
                print('something went wrong for', model, e)
                errors += 1
    print(f"\nModel {model} created {len(sequences)} ({errors=}) unique sequences:")
    print(json.dumps(list(sequences)))
    results.append((len(sequences), model))
# Testing completion models
for model in completion_models:
    sequences = set()
    errors = 0
    with TimeIt(model):
        for _ in range(C):
            try:
                completion = openai.Completion.create(
                    model=model,
                    prompt=prompt,
                    max_tokens=N,
                    temperature=0,
                    logit_bias = {"50256": -100.0}, # prevent EOS
                )
                sequences.add(completion.choices[0].text)
                sleep(1)
            except Exception as e:
                print('something went wrong for', model, e)
                errors += 1
    print(f"\nModel {model} created {len(sequences)} ({errors=}) unique sequences:")
    print(json.dumps(list(sequences)))
    results.append((len(sequences), model))

# Printing table of results
print("\nTable of Results:")
print("Num_Sequences\tModel_Name")
for num_sequences, model_name in results:
    print(f"{num_sequences}\t{model_name}")