gpt.py

import os
import time
import traceback
from pprint import pprint

from celery import Celery
import openai
from util.util import retry, timestamp
from util.gpt_util import parse_logit_bias, parse_stop, get_correct_key
import requests
import codecs
import json

# response dictionary type
'''
{
    "completions": [{'text': string
                     'tokens': [token_data]
                     'finishReason': string}]
    "prompt": {
                'text': string,
                ? 'tokens': [token_data]
              }
    "id": string
    "model": string
    "timestamp": timestamp
}
'''

# token data dictionary type
'''
{
    'generatedToken': {'logprob': float,
                       'token': string}
    'position': {'end': int, 'start': int}
    ? 'counterfactuals': [{'token': float)}]
}
'''

# finishReason
'''
"finishReason": {"reason": "stop" | "length",
                 ? "sequence": string }
'''



#ai21_api_key = os.environ.get("AI21_API_KEY", None)

client = openai.Client(api_key=os.environ.get("OPENAI_API_KEY", 'placeholder'))


def gen(prompt, settings, config, **kwargs):
    if settings["stop"]:
        stop = parse_stop(settings["stop"])
    else:
        stop = None
    if settings["logit_bias"]:
        logit_bias = parse_logit_bias(settings["logit_bias"])
    else:
        logit_bias = None
    #if config['OPENAI_API_KEY']:
    model_info = config['models'][settings['model']]
    # print('model info:', model_info)
    client.base_url = model_info['api_base'] if model_info['api_base'] else "https://api.openai.com/v1"

    ai21_api_key = kwargs.get('AI21_API_KEY', None)
    ai21_api_key = ai21_api_key if ai21_api_key else os.environ.get("AI21_API_KEY", None)
    client.api_key, client.organization = get_correct_key(model_info['type'], kwargs)

    # print('openai api base: ' + openai.api_base)

    # print('openai api key: ' + openai.api_key)

    # if config['AI21_API_KEY']:
        #TODO
        # ai21_api_key = config['AI21_API_KEY']
    try:
        response, error = generate(prompt=prompt,
                                length=settings['response_length'],
                                num_continuations=settings['num_continuations'],
                                temperature=settings['temperature'],
                                logprobs=settings['logprobs'],
                                top_p=settings['top_p'],
                                model=settings['model'],
                                stop=stop,
                                logit_bias=logit_bias,
                                config=config,
                                ai21_api_key=ai21_api_key,
                                )
        return response, error
    except Exception as e:
        print(e)
        return None, e


def generate(config, **kwargs):
    #pprint(kwargs)
    model_type = config['models'][kwargs['model']]['type']
    if model_type == 'ai21':
        response, error = ai21_generate(api_key=kwargs['ai21_api_key'], **kwargs)#config['AI21_API_KEY'], **kwargs)
        #save_response_json(response.json(), 'examples/AI21_response.json')
        if not error:
            formatted_response = format_ai21_response(response.json(), model=kwargs['model'])
            #save_response_json(formatted_response, 'examples/AI21_formatted_response.json')
            return formatted_response, error
        else:
            return response, error
    elif model_type in ('openai', 'openai-custom', 'gooseai', 'openai-chat', 'together', 'llama-cpp'):
        is_chat = model_type in ('openai-chat',)
        # for some reason, Together AI ignores the echo parameter
        echo = model_type not in ('together', 'openai-chat')
        # TODO: Together AI and chat inference breaks if logprobs is set to 0
        assert kwargs['logprobs'] > 0 or model_type not in ('together',), \
            "Logprobs must be greater than 0 for model type Together AI"
        # llama-cpp-python doesn't support batched inference yet: https://github.com/abetlen/llama-cpp-python/issues/771
        needs_multiple_calls = model_type in ('llama-cpp',)
        if needs_multiple_calls:
            required_calls = kwargs['num_continuations']
            kwargs['num_continuations'] = 1
            responses = []
            for _ in range(required_calls):
                response, error = openAI_generate(model_type, **kwargs)
                responses.append(response)
            response = responses[-1]
            response['choices'] = [r['choices'][0] for r in responses]
        else:
            # TODO OpenAI errors
            response, error = openAI_generate(model_type, **kwargs)
        #save_response_json(response, 'examples/openAI_response.json')
        formatted_response = format_openAI_response(response, kwargs['prompt'], echo=echo, is_chat=is_chat)
        #save_response_json(formatted_response, 'examples/openAI_formatted_response.json')
        return formatted_response, error


def completions_text(response):
    return [completion['text'] for completion in response['completions']]


def save_response_json(response, filename):
    with open(filename, 'w') as f:
        json.dump(response, f)

#################################
#   Janus
#################################

redis_url = os.environ.get("JANUS_REDIS", None)
app = Celery(
    # 'janus',
    broker=redis_url,
    backend=redis_url,
)

# get_gpt_response(prompt, memory, retry=True) -> result, error
janus_task = "janus.my_celery.tasks.get_gpt_response"


def janus_generate(prompt, memory=""):
    assert isinstance(prompt, str) and isinstance(memory, str)
    celery_task = app.send_task(janus_task, args=[prompt, memory])
    print("Sent to janus")
    result, error = celery_task.get()
    return result, error


#################################
#   OpenAI
#################################

#openai.api_key = os.environ.get("OPENAI_API_KEY", None)


def fix_openAI_token(token):
    # if token is a byte string, convert to string
    # TODO this doesn't work
    decoded = codecs.decode(token, "unicode-escape")
    return decoded
    # byte_token = decoded.encode('raw_unicode_escape')
    # return byte_token.decode('utf-8')


def format_openAI_token_dict(completion, token, i, offset):
    calculated_offset = len(token) + offset
    token_dict = {'generatedToken': {'token': token,
                                     'logprob': completion['logprobs']['token_logprobs'][i]},
                  'position': calculated_offset}
    if completion['logprobs'].get('top_logprobs', None) is not None and \
        completion['logprobs']['top_logprobs']:
        openai_counterfactuals = completion['logprobs']['top_logprobs'][i]
        if openai_counterfactuals:
            sorted_counterfactuals = {k: v for k, v in
                                    sorted(openai_counterfactuals.items(), key=lambda item: item[1], reverse=True)}
            token_dict['counterfactuals'] = sorted_counterfactuals
    else:
        token_dict['counterfactuals'] = None
    return token_dict, calculated_offset

def format_openAI_chat_token_dict(content_token, i):
    token_dict = {
        'generatedToken': {'token': content_token['token'],
                           'logprob': content_token['logprob']},
                           'position': i,
                           'counterfactuals' : {c['token']: c['logprob'] for c in content_token['top_logprobs']}
        }
    return token_dict

def format_openAI_completion(completion, prompt_offset, prompt_end_index, is_chat):
    if 'text' in completion:
        completion_text = completion['text']
    else:
        completion_text = completion['message']['content']
    completion_dict = {'text': completion_text[prompt_offset:],
                       'finishReason': completion['finish_reason'],
                       'tokens': []}
    offset = prompt_offset
    if is_chat:
        for i, token in enumerate(completion['logprobs']['content']):
            token_dict = format_openAI_chat_token_dict(token, i)
            completion_dict['tokens'].append(token_dict)
    else:
        for i, token in enumerate(completion['logprobs']['tokens'][prompt_end_index:]):
            j = i + prompt_end_index
            token_dict, offset = format_openAI_token_dict(completion, token, j, offset)
            completion_dict['tokens'].append(token_dict)
    return completion_dict


def format_openAI_prompt(completion, prompt, prompt_end_index):
    prompt_dict = {'text': prompt, 'tokens': []}
    # loop over tokens until offset >= prompt length
    offset = 0
    for i, token in enumerate(completion['logprobs']['tokens'][:prompt_end_index]):
        token_dict, offset = format_openAI_token_dict(completion, token, i, offset)
        prompt_dict['tokens'].append(token_dict)

    return prompt_dict


def format_openAI_response(response, prompt, echo, is_chat):
    if echo:
        prompt_end_index = response['usage']['prompt_tokens']
        prompt_dict = format_openAI_prompt(response['choices'][0],
                                                             prompt,
                                                             prompt_end_index)
    else:
        prompt_dict = {'text': prompt, 'tokens': None}
        prompt_end_index = 0
        #prompt = ''

    prompt_offset = len(prompt) if echo else 0

    response_dict = {'completions': [format_openAI_completion(completion, prompt_offset, prompt_end_index, is_chat) for
                                     completion in response['choices']],
                     'prompt': prompt_dict,
                     'id': response['id'],
                     'model': response['model'],
                     'timestamp': timestamp()}
    return response_dict


@retry(n_tries=3, delay=1, backoff=2, on_failure=lambda *args, **kwargs: ("", None))
def openAI_generate(model_type, prompt, length=150, num_continuations=1, logprobs=10, temperature=0.8, top_p=1, stop=None,
                    model='davinci', logit_bias=None, **kwargs):
    if not logit_bias:
        logit_bias = {}
    params = {
        'temperature': temperature,
        'max_tokens': length,
        'top_p': top_p,
        'logprobs': logprobs,
        'logit_bias': logit_bias,
        'n': num_continuations,
        'stop': stop,
        'model': model,
        #**kwargs
    }
    if model_type == 'openai-chat':
        params['messages'] = [{ 'role': "assistant", 'content': prompt }]
        params['logprobs'] = True
        params['top_logprobs'] = logprobs
        response = client.chat.completions.create(
            **params
        ).to_dict()
    else:
        params['prompt'] = prompt
        params['echo'] = True
        response = client.completions.create(
            **params
        ).to_dict()

    return response, None


def search(query, documents, engine="curie"):
    return client.Engine(engine).search(
        documents=documents,
        query=query
    )


#################################
#   AI21
#################################


def fix_ai21_tokens(token):
    return token.replace("▁", " ").replace("<|newline|>", "\n")

def ai21_token_position(textRange, text_offset):
    return {'start': textRange['start'] + text_offset,
            'end': textRange['end'] + text_offset}

def format_ai21_token_data(token, prompt_offset=0):
    token_dict = {'generatedToken': {'token': fix_ai21_tokens(token['generatedToken']['token']),
                                     'logprob': token['generatedToken']['logprob']},
                  'position': ai21_token_position(token['textRange'], prompt_offset)}
    if token['topTokens']:
        token_dict['counterfactuals'] = {fix_ai21_tokens(c['token']): c['logprob'] for c in token['topTokens']}
    else:
        token_dict['counterfactuals'] = None
    return token_dict


def format_ai21_completion(completion, prompt_offset=0):
    completion_dict = {'text': completion['data']['text'],
                       'tokens': [format_ai21_token_data(token, prompt_offset) for token in completion['data']['tokens']],
                       'finishReason': completion['finishReason']['reason']}
    return completion_dict


def format_ai21_response(response, model):
    prompt = response['prompt']['text']
    response_dict = {'completions': [format_ai21_completion(completion, prompt_offset=len(prompt)) for completion in response['completions']],
                     'prompt': {'text': prompt,
                                'tokens': [format_ai21_token_data(token, prompt_offset=0) for token in response['prompt']['tokens']]},
                     'id': response['id'],
                     'model': model,
                     'timestamp': timestamp()}
    return response_dict


def ai21_generate(prompt, length=150, num_continuations=1, logprobs=10, temperature=0.8, top_p=1, stop=None,
                  engine='j1-large', api_key=None, **kwargs):
    stop = stop if stop else []
    request_json = {
        "prompt": prompt,
        "numResults": num_continuations,
        "maxTokens": length,
        "stopSequences": stop,
        "topKReturn": logprobs,
        "temperature": temperature,
        "topP": top_p,
    }
    try:
        response = requests.post(
            f"https://api.ai21.com/studio/v1/{engine}/complete",
            headers={"Authorization": f"Bearer {api_key}"},
            json=request_json,
        )
    except requests.exceptions.ConnectionError:
        return None, 'Connection error'
    error = None
    if response.status_code != 200:
        error = f'Bad status code {response.status_code}'
        print(request_json)
    return response, error


if __name__ == "__main__":
    pass