-
Notifications
You must be signed in to change notification settings - Fork 0
/
openai_helper.py
56 lines (45 loc) · 1.79 KB
/
openai_helper.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
import openai
from secret_key import openai_key
import json
import pandas as pd
openai.api_key = openai_key
def extract_financial_data(text):
prompt = get_prompt_financial() + text
response = openai.ChatCompletion.create(
model="gpt-3.5-turbo",
messages=[{"role": "user","content": prompt}]
)
content = response.choices[0]['message']['content']
try:
data = json.loads(content)
return pd.DataFrame(data.items(), columns=["Measure", "Value"])
except (json.JSONDecodeError, IndexError):
pass
return pd.DataFrame({
"Measure": ["Company Name", "Stock Symbol", "Revenue", "Net Income", "EPS"],
"Value": ["", "", "", "", ""]
})
def get_prompt_financial():
return '''Please retrieve company name, revenue, net income and earnings per share (a.k.a. EPS)
from the following news article. If you can't find the information from this article
then return "". Do not make things up.
Then retrieve a stock symbol corresponding to that company. For this you can use
your general knowledge (it doesn't have to be from this article). Always return your
response as a valid JSON string. The format of that string should be this,
{
"Company Name": "Walmart",
"Stock Symbol": "WMT",
"Revenue": "12.34 million",
"Net Income": "34.78 million",
"EPS": "2.1 $"
}
News Article:
============
'''
# Press the green button in the gutter to run the script.
if __name__ == '__main__':
text = '''
Tesla's Earning news in text format: Tesla's earning this quarter blew all the estimates. They reported 4.5 billion $ profit against a revenue of 30 billion $. Their earnings per share was 2.3 $
'''
df = extract_financial_data(text)
print(df.to_string())