RelationalIntentionGraph/prompting.py at main · HKUST-KnowComp/RelationalIntentionGraph · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
import openai
from util import get_message, get_simple_message
import sys
sys.path.append(r'./')
openai.api_type = "azure"
openai.api_base = "https://hkust.azure-api.net"
openai.api_key = "6c25a1df10384b90afedc13130074359"
openai.api_version = "2023-05-15"
from tqdm import tqdm
from tenacity import retry,stop_after_attempt,stop_after_delay,wait_fixed
import json


# retry 10 times, interval 10s, waiting 100s
@retry(stop=(stop_after_delay(100) | stop_after_attempt(10)), wait=wait_fixed(10))
def get_message_from_api(engine_name, message):

    assert engine_name in ["gpt-35-turbo", "gpt-4"]
    response = openai.ChatCompletion.create(
            engine = engine_name,
            temperature = 0,
            messages = [{"role":"user","content":message}])
    a = response.get("choices")[0]["message"]["content"]

    usage = response.get("usage")
    completion_tokens = usage["completion_tokens"]
    prompt_tokens = usage["prompt_tokens"]

    return a, prompt_tokens, completion_tokens

def get_gpt_answer(engine_name, sample_session_path, sample_item_path):


    generation_result_path = f'./generation_results/{engine_name}_answer_test.json'
    f = open(generation_result_path,'w',encoding='utf-8')
    idx = 0
    for sessions, message in tqdm(get_simple_message(sample_session_path,sample_item_path)):
        try:
            a, prompt_tokens, completion_tokens = get_message_from_api(engine_name, message=message)

            result_dict = {
                "Session": sessions,
                "Message": message,
                "Answer": a,
                "Prompt_tokens": prompt_tokens,
                "Completion_tokens": completion_tokens,
            }

            result_string = json.dumps(result_dict, ensure_ascii=False)
            f.write(result_string+'\n')


        except Exception as e:
            print(f'session {idx} have expection:')
            print(e)
            a = 'ERROR:'+str(idx)


        idx+=1

        if idx%10==0:
            f.flush()


if __name__ == "__main__":


    sample_session_path = r'./data/uk_sample_1000.txt'
    sample_item_path= r'./data/uk_item_sample_1000.txt'

    # for engine_name in ["gpt-35-turbo", "gpt-4"]:
    #     get_gpt_answer(engine_name, sample_session_path, sample_item_path)

    for engine_name in ["gpt-4"]:
        get_gpt_answer(engine_name, sample_session_path, sample_item_path)

    # balance before gpt4: HK$ 999.997708
    # balance after: HK$ 857.83394