-
Notifications
You must be signed in to change notification settings - Fork 0
/
assistant.py
67 lines (56 loc) · 2.19 KB
/
assistant.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
import chromadb
from openai import OpenAI
import os
import tiktoken
try:
chroma_client = chromadb.PersistentClient("data/chroma")
collection = chroma_client.get_collection(name="fine_food_reviews_1k")
except ValueError:
print("No collection to retrieve from!")
exit()
# models
embedding_model = "text-embedding-3-small"
embedding_encoding = "cl100k_base"
gpt_model = "gpt-3.5-turbo"
token_budget = 4096 - 500
openai_client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
print("(Optional) Enter product id: ")
product_id = input().strip()
print("(Optional) Enter user id: ")
user_id = input().strip()
print("Enter your prompt:")
query_text = input().strip()
metafields = {}
if product_id: metafields["ProductId"] = product_id
if user_id: metafields["UserId"] = user_id
query_embedding = openai_client.embeddings.create(input = [query_text], model=embedding_model).data[0].embedding
results = collection.query(query_embeddings=[query_embedding], where=metafields if metafields else None)
reviews = collection.get(ids=results["ids"][0], include=["documents"])["documents"]
def num_tokens(text, model):
"""Return the number of tokens in a string."""
encoding = tiktoken.encoding_for_model(model)
return len(encoding.encode(text))
introduction = "Use the following retrieved reviews from the dataset to answer the subsequent prompt."
prompt = f"\n\nPrompt: {query_text}"
message = introduction
for i, review in enumerate(reviews):
next_review = f'\n\nReview #{i+1}:\n"""\n{review}\n"""'
if (num_tokens(message + next_review + prompt, model=gpt_model) > token_budget):
break
else:
message += next_review
message += prompt
messages = [
{"role": "system", "content": "You answer questions about the Amazon Fine Food Reviews dataset, which consists of reviews of fine foods from Amazon."},
{"role": "user", "content": message},
]
response = openai_client.chat.completions.create(
model=gpt_model,
messages=messages,
temperature=0
)
print("\nMessage sent to assistant:\n# # # # # # # # # #\n")
print(message)
print("\n# # # # # # # # # #\n")
response_message = response.choices[0].message.content
print(f"Assistant response ({gpt_model}): \n{response_message}")