-
Notifications
You must be signed in to change notification settings - Fork 418
/
Copy pathinference.py
executable file
·188 lines (181 loc) · 8.83 KB
/
inference.py
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
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
import time, tiktoken
from openai import OpenAI
import openai
import os, anthropic, json
TOKENS_IN = dict()
TOKENS_OUT = dict()
encoding = tiktoken.get_encoding("cl100k_base")
def curr_cost_est():
costmap_in = {
"gpt-4o": 2.50 / 1000000,
"gpt-4o-mini": 0.150 / 1000000,
"o1-preview": 15.00 / 1000000,
"o1-mini": 3.00 / 1000000,
"claude-3-5-sonnet": 3.00 / 1000000,
"deepseek-chat": 1.00 / 1000000,
"o1": 15.00 / 1000000,
}
costmap_out = {
"gpt-4o": 10.00/ 1000000,
"gpt-4o-mini": 0.6 / 1000000,
"o1-preview": 60.00 / 1000000,
"o1-mini": 12.00 / 1000000,
"claude-3-5-sonnet": 12.00 / 1000000,
"deepseek-chat": 5.00 / 1000000,
"o1": 60.00 / 1000000,
}
return sum([costmap_in[_]*TOKENS_IN[_] for _ in TOKENS_IN]) + sum([costmap_out[_]*TOKENS_OUT[_] for _ in TOKENS_OUT])
def query_model(model_str, prompt, system_prompt, openai_api_key=None, anthropic_api_key=None, tries=5, timeout=5.0, temp=None, print_cost=True, version="1.5"):
preloaded_api = os.getenv('OPENAI_API_KEY')
if openai_api_key is None and preloaded_api is not None:
openai_api_key = preloaded_api
if openai_api_key is None and anthropic_api_key is None:
raise Exception("No API key provided in query_model function")
if openai_api_key is not None:
openai.api_key = openai_api_key
os.environ["OPENAI_API_KEY"] = openai_api_key
if anthropic_api_key is not None:
os.environ["ANTHROPIC_API_KEY"] = anthropic_api_key
for _ in range(tries):
try:
if model_str == "gpt-4o-mini" or model_str == "gpt4omini" or model_str == "gpt-4omini" or model_str == "gpt4o-mini":
model_str = "gpt-4o-mini"
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": prompt}]
if version == "0.28":
if temp is None:
completion = openai.ChatCompletion.create(
model=f"{model_str}", # engine = "deployment_name".
messages=messages
)
else:
completion = openai.ChatCompletion.create(
model=f"{model_str}", # engine = "deployment_name".
messages=messages, temperature=temp
)
else:
client = OpenAI()
if temp is None:
completion = client.chat.completions.create(
model="gpt-4o-mini-2024-07-18", messages=messages, )
else:
completion = client.chat.completions.create(
model="gpt-4o-mini-2024-07-18", messages=messages, temperature=temp)
answer = completion.choices[0].message.content
elif model_str == "claude-3.5-sonnet":
client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
message = client.messages.create(
model="claude-3-5-sonnet-latest",
system=system_prompt,
messages=[{"role": "user", "content": prompt}])
answer = json.loads(message.to_json())["content"][0]["text"]
elif model_str == "gpt4o" or model_str == "gpt-4o":
model_str = "gpt-4o"
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": prompt}]
if version == "0.28":
if temp is None:
completion = openai.ChatCompletion.create(
model=f"{model_str}", # engine = "deployment_name".
messages=messages
)
else:
completion = openai.ChatCompletion.create(
model=f"{model_str}", # engine = "deployment_name".
messages=messages, temperature=temp)
else:
client = OpenAI()
if temp is None:
completion = client.chat.completions.create(
model="gpt-4o-2024-08-06", messages=messages, )
else:
completion = client.chat.completions.create(
model="gpt-4o-2024-08-06", messages=messages, temperature=temp)
answer = completion.choices[0].message.content
elif model_str == "deepseek-chat":
model_str = "deepseek-chat"
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": prompt}]
if version == "0.28":
raise Exception("Please upgrade your OpenAI version to use DeepSeek client")
else:
deepseek_client = OpenAI(
api_key=os.getenv('DEEPSEEK_API_KEY'),
base_url="https://api.deepseek.com/v1"
)
if temp is None:
completion = deepseek_client.chat.completions.create(
model="deepseek-chat",
messages=messages)
else:
completion = deepseek_client.chat.completions.create(
model="deepseek-chat",
messages=messages,
temperature=temp)
answer = completion.choices[0].message.content
elif model_str == "o1-mini":
model_str = "o1-mini"
messages = [
{"role": "user", "content": system_prompt + prompt}]
if version == "0.28":
completion = openai.ChatCompletion.create(
model=f"{model_str}", # engine = "deployment_name".
messages=messages)
else:
client = OpenAI()
completion = client.chat.completions.create(
model="o1-mini-2024-09-12", messages=messages)
answer = completion.choices[0].message.content
elif model_str == "o1":
model_str = "o1"
messages = [
{"role": "user", "content": system_prompt + prompt}]
if version == "0.28":
completion = openai.ChatCompletion.create(
model="o1-2024-12-17", # engine = "deployment_name".
messages=messages)
else:
client = OpenAI()
completion = client.chat.completions.create(
model="o1-2024-12-17", messages=messages)
answer = completion.choices[0].message.content
elif model_str == "o1-preview":
model_str = "o1-preview"
messages = [
{"role": "user", "content": system_prompt + prompt}]
if version == "0.28":
completion = openai.ChatCompletion.create(
model=f"{model_str}", # engine = "deployment_name".
messages=messages)
else:
client = OpenAI()
completion = client.chat.completions.create(
model="o1-preview", messages=messages)
answer = completion.choices[0].message.content
try:
if model_str in ["o1-preview", "o1-mini", "claude-3.5-sonnet", "o1"]:
encoding = tiktoken.encoding_for_model("gpt-4o")
elif model_str in ["deepseek-chat"]:
encoding = tiktoken.encoding_for_model("cl100k_base")
else:
encoding = tiktoken.encoding_for_model(model_str)
if model_str not in TOKENS_IN:
TOKENS_IN[model_str] = 0
TOKENS_OUT[model_str] = 0
TOKENS_IN[model_str] += len(encoding.encode(system_prompt + prompt))
TOKENS_OUT[model_str] += len(encoding.encode(answer))
if print_cost:
print(f"Current experiment cost = ${curr_cost_est()}, ** Approximate values, may not reflect true cost")
except Exception as e:
if print_cost:
print(f"Cost approximation has an error? {e}")
return answer
except Exception as e:
print("Inference Exception:", e)
time.sleep(timeout)
continue
raise Exception("Max retries: timeout")
#print(query_model(model_str="o1-mini", prompt="hi", system_prompt="hey"))