openai_llm_main.py

import os
import sys
import copy
import json

import time
import ray
import openai
from tqdm import tqdm
from rich.console import Console
import concurrent.futures

from dataset import PhotoChatDataset
from config import ex
from utils import fixed_seed
from load_dataset import prepare_prompt_dataset


console = Console()
error_console = Console(stderr=True, style='bold red')


mode_to_api_caller = {
    'chat': openai.ChatCompletion,
    'completion': openai.Completion,
}


def call_gpt3(inputs, mode, model_name, params):
    # call GPT-3 API until result is provided and then return it
    response = None
    received = False
    
    openai.organization = os.getenv("OPENAI_ORGANIZATION_ID")
    openai.api_key = os.getenv("OPENAI_API_KEY")


    while not received:
        try:
            if mode == 'chat':
                response = mode_to_api_caller[mode].create(model=model_name, messages=inputs, **params)
            elif mode == 'completion':
                response = mode_to_api_caller[mode].create(model=model_name, prompt=inputs, **params)
            received = True
        except openai.error.OpenAIError as e:
            print(f"OpenAIError: {e}.")
            error = sys.exc_info()[0]
            #if error == openai.error.InvalidRequestError:
                # something is wrong: e.g., prompt too long
            #    print(f"InvalidRequestError\nPrompt passed in:\n\n{inputs}\n\n")
            #    assert False
            
            time.sleep(2)
    
    return response

def openai_inference(instance, config):

    if config['mode'] == 'chat':
        prompt = [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": instance['prompt']}
        ]
    elif config['mode'] == 'completion':
        prompt = instance['prompt']

    resp = call_gpt3(prompt, config['mode'], config['model']['name'], config['openai_params'])

    #instance["{}_generation".format(config["task_num"])] = resp.choices[0].message["content"]
    
    #return instance
    return resp.choices[0].message["content"]

def multi_run_chat(dataset, config, result_save_dir):
    outputs = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor:
        futures = []
        
        for instance in tqdm(dataset, total=len(dataset)):
            copied_instance = copy.deepcopy(instance)
            
            future = executor.submit(openai_inference, copied_instance, config)
            futures.append(future)
        
        for future in tqdm(concurrent.futures.as_completed(futures), total=len(futures)):
            ret = future.result()
            outputs.append(ret[0])

def run_chat(dataset, config):
    outputs = []
    for instance in tqdm(dataset, total=len(dataset)):
        output = openai_inference(instance, config)
        copied_instance = copy.deepcopy(instance)
        copied_instance['{}_generation'.format(config['task_num'])] = output
        outputs.append(copied_instance)
    
    assert len(outputs) == len(dataset)
    return outputs

@ex.automain
def main(_config):
    _config = copy.deepcopy(_config)
    
    # set CUDA
    if isinstance(_config["gpu_id"], list):
        cuda_devices = ",".join([
            str(ele) for ele in _config["gpu_id"]
        ])
    else:
        cuda_devices = str(_config["gpu_id"])
    os.environ["CUDA_VISIBLE_DEVICES"] = cuda_devices
    console.log("GPU: {}".format(cuda_devices))

    # set seed
    fixed_seed(_config["seed"])

    result_save_dir = os.path.join(
        _config['result_save_dir'], 
        _config['file_version'],
        _config['dataset_name'], 
        _config['template_name'],
        'seed_{}'.format(_config['seed'])
    )
    os.makedirs(result_save_dir, exist_ok=True)
    console.log(f"Result save directory: {result_save_dir}")

    dataset = prepare_prompt_dataset(_config)

    #print(dataset[0])
    #print(dataset[0]['prompt'])

    import random
    random.shuffle(dataset)

    sampled_dataset = dataset[:12]
    save_str = []
    for ele in sampled_dataset:
        save_str.append(ele['prompt'] + ' ' + ele['label'])

    with open('{}.txt'.format(_config['template_name']), 'w') as f:
        f.write('\n\n\n'.join(save_str))

    assert False

    if _config["do_sample"]:
        dataset = dataset[:_config["sample_num"]]
        console.log("Sampling is applied for the fast debug.")

    mode = _config['mode']
    if mode == 'chat':
        llm_outputs = run_chat(dataset, _config)
    elif mode == 'completion':
        run_completion()


    if _config['do_cot']:
        with open(os.path.join(result_save_dir, 'cot-{}.json'.format(_config['model']['id'])), 'w') as f:
            json.dump(llm_outputs, f, ensure_ascii=False, indent='\t')

    else:
        with open(os.path.join(result_save_dir, '{}.json'.format(_config['model']['id'])), 'w') as f:
            json.dump(llm_outputs, f, ensure_ascii=False, indent='\t')