Spaces:
Runtime error
Runtime error
军舰 commited on
Commit ·
456be22
1
Parent(s): e0c4e40
add user set OpenAI API Key
Browse files
chat.py
ADDED
|
@@ -0,0 +1,75 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import openai
|
| 2 |
+
import tiktoken
|
| 3 |
+
|
| 4 |
+
|
| 5 |
+
class Conversation:
|
| 6 |
+
def __init__(self, prompt, model="gpt-3.5-turbo", temperature=0.8, max_tokens=250):
|
| 7 |
+
self.prompt = prompt
|
| 8 |
+
self.model = model
|
| 9 |
+
self.temperature = temperature
|
| 10 |
+
self.max_tokens = max_tokens
|
| 11 |
+
|
| 12 |
+
self._init_messages()
|
| 13 |
+
|
| 14 |
+
def _init_messages(self):
|
| 15 |
+
self.messages = [{"role": "system", "content": self.prompt}]
|
| 16 |
+
|
| 17 |
+
def reset(self):
|
| 18 |
+
self._init_messages()
|
| 19 |
+
|
| 20 |
+
def ask(self, question, pprint=True):
|
| 21 |
+
self.messages.append({"role": "user", "content": question})
|
| 22 |
+
|
| 23 |
+
if self.num_tokens(self.messages, self.model) >= self.max_tokens:
|
| 24 |
+
if len(self.messages) > 3:
|
| 25 |
+
self.messages = self.messages[:1] + self.messages[3:] # remove the first user message
|
| 26 |
+
else:
|
| 27 |
+
return "Error: max tokens exceeded."
|
| 28 |
+
|
| 29 |
+
try:
|
| 30 |
+
response = openai.ChatCompletion.create(
|
| 31 |
+
model=self.model,
|
| 32 |
+
messages=self.messages
|
| 33 |
+
)
|
| 34 |
+
except Exception as e:
|
| 35 |
+
return e
|
| 36 |
+
|
| 37 |
+
if pprint:
|
| 38 |
+
print(f"tiktoken: {self.num_tokens(self.messages, self.model)}\ntokens: {response['usage']}")
|
| 39 |
+
|
| 40 |
+
assistant_message = response["choices"][0]["message"]["content"]
|
| 41 |
+
self.messages.append({"role": "assistant", "content": assistant_message})
|
| 42 |
+
|
| 43 |
+
return assistant_message
|
| 44 |
+
|
| 45 |
+
def num_tokens(self, messages, model):
|
| 46 |
+
"""Returns the number of tokens used by a list of messages."""
|
| 47 |
+
try:
|
| 48 |
+
encoding = tiktoken.encoding_for_model(model)
|
| 49 |
+
except KeyError:
|
| 50 |
+
print("Warning: model not found. Using cl100k_base encoding.")
|
| 51 |
+
encoding = tiktoken.get_encoding("cl100k_base")
|
| 52 |
+
if model == "gpt-3.5-turbo":
|
| 53 |
+
print("Warning: gpt-3.5-turbo may change over time. Returning num tokens assuming gpt-3.5-turbo-0301.")
|
| 54 |
+
return self.num_tokens(messages, model="gpt-3.5-turbo-0301")
|
| 55 |
+
elif model == "gpt-4":
|
| 56 |
+
print("Warning: gpt-4 may change over time. Returning num tokens assuming gpt-4-0314.")
|
| 57 |
+
return self.num_tokens(messages, model="gpt-4-0314")
|
| 58 |
+
elif model == "gpt-3.5-turbo-0301":
|
| 59 |
+
tokens_per_message = 4 # every message follows <|start|>{role/name}\n{content}<|end|>\n
|
| 60 |
+
tokens_per_name = -1 # if there's a name, the role is omitted
|
| 61 |
+
elif model == "gpt-4-0314":
|
| 62 |
+
tokens_per_message = 3
|
| 63 |
+
tokens_per_name = 1
|
| 64 |
+
else:
|
| 65 |
+
raise NotImplementedError(f"""num_tokens_from_messages() is not implemented for model {model}. See https://github.com/openai/openai-python/blob/main/chatml.md for information on how messages are converted to tokens.""")
|
| 66 |
+
num_tokens = 0
|
| 67 |
+
for message in messages:
|
| 68 |
+
num_tokens += tokens_per_message
|
| 69 |
+
for key, value in message.items():
|
| 70 |
+
num_tokens += len(encoding.encode(value))
|
| 71 |
+
if key == "name":
|
| 72 |
+
num_tokens += tokens_per_name
|
| 73 |
+
num_tokens += 3 # every reply is primed with <|start|>assistant<|message|>
|
| 74 |
+
return num_tokens
|
| 75 |
+
|