Larfii commited on
Commit
a8a6171
·
1 Parent(s): ddb02a2
Files changed (1) hide show
  1. examples/generate_query.py +56 -0
examples/generate_query.py ADDED
@@ -0,0 +1,56 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+
3
+ from openai import OpenAI
4
+
5
+ os.environ["OPENAI_API_KEY"] = ""
6
+
7
+ def openai_complete_if_cache(
8
+ model="gpt-4o-mini", prompt=None, system_prompt=None, history_messages=[], **kwargs
9
+ ) -> str:
10
+ openai_client = OpenAI()
11
+
12
+ messages = []
13
+ if system_prompt:
14
+ messages.append({"role": "system", "content": system_prompt})
15
+ messages.extend(history_messages)
16
+ messages.append({"role": "user", "content": prompt})
17
+
18
+ response = openai_client.chat.completions.create(
19
+ model=model, messages=messages, **kwargs
20
+ )
21
+ return response.choices[0].message.content
22
+
23
+
24
+ if __name__ == "__main__":
25
+ description = ""
26
+ prompt = f"""
27
+ Given the following description of a dataset:
28
+
29
+ {description}
30
+
31
+ Please identify 5 potential users who would engage with this dataset. For each user, list 5 tasks they would perform with this dataset. Then, for each (user, task) combination, generate 5 questions that require a high-level understanding of the entire dataset.
32
+
33
+ Output the results in the following structure:
34
+ - User 1: [user description]
35
+ - Task 1: [task description]
36
+ - Question 1:
37
+ - Question 2:
38
+ - Question 3:
39
+ - Question 4:
40
+ - Question 5:
41
+ - Task 2: [task description]
42
+ ...
43
+ - Task 5: [task description]
44
+ - User 2: [user description]
45
+ ...
46
+ - User 5: [user description]
47
+ ...
48
+ """
49
+
50
+ result = openai_complete_if_cache(model='gpt-4o-mini', prompt=prompt)
51
+
52
+ file_path = f"./queries.txt"
53
+ with open(file_path, "w") as file:
54
+ file.write(result)
55
+
56
+ print(f"Queries written to {file_path}")