update streamlit app

.gitignore

@@ -1,4 +1,5 @@
 *.idea
 prompt.txt
 key.txt
-*.ipynb_checkpoints/*
+*.ipynb_checkpoints/*
+*.csv

flask_app.py

@@ -62,7 +62,7 @@ def intakeMessage(msg):
 
 
     print("output:", output)
-    if type(output) == pd.core.series.Series and len(output) == 1:
+    if dtype == pd.core.series.Series and len(output) == 1:
         command = f"output = {response.choices[0].text.replace('A: ', '')}.iloc[0]"
         exec(command)
         output = int(output)

streamlit_app.py

@@ -1,33 +1,85 @@
 import streamlit as st
 import pandas as pd
-import time
+import openai
+import numpy as np
 
 
-df = pd.DataFrame(
-    [
-        {'building': 'Empire State', 'location': "NYC", 'height':1000},
-        {'building': 'B', 'location': "LA", 'height':600},
-        {'building': 'c', 'location': "Paris", 'height':900},
-        {'building': 'C', 'location': "SF", 'height':800},
-        {'building': 'P', 'location': "Dallas", 'height':500},
-    ]
-)
+def file_reader(file):
+    file_type = file.name.split('.')[-1]
 
+    # TODO add more file types
+    if file_type == 'csv':
+        return pd.read_csv(file)
+
+    if file_type == 'parquet':
+        return pd.read_parquet(file_type)
+
+    return None
+
+
+f = open("key.txt", "r")
+key = f.read()
+f.close()
+f = open("prompt.txt", "r")
+prompt = f.read()
+f.close()
+openai.api_key = key
 
 
 st.title('IZE')
 st.subheader('Use natural language to visualize and analyze your data')
 
-st.dataframe(df.head())
+uploaded_file = st.file_uploader("Choose a file to analyze")
+
+
+if uploaded_file is not None:
+    df = file_reader(uploaded_file)
+
+    if df is None:
+        st.write(f'{uploaded_file.name.split(".")[-1]} is not a supported file type')
+
+
+if uploaded_file is not None and df is not None:
+
+    st.dataframe(df.head())
+
+    example = 'Ex: How many rows are there?'
+    user_input = st.text_input("Input your Query here", example)
+
+    if user_input and user_input != example:
+        with st.spinner('Loading...'):
+            new_prompt = f"{prompt}\n\nQ:{user_input} columns: {list(df.columns)}\n"
+
+            response = openai.Completion.create(engine='davinci',
+                                                prompt=new_prompt,
+                                                stop='\n',
+                                                temperature=0,
+                                                top_p=1,
+                                                frequency_penalty=0,
+                                                presence_penalty=0,
+                                                max_tokens=150
+                                                )
+
+            output = ''
+            command = f"output = {response.choices[0].text.replace('A: ', '')}; dtype = type(output)"
+            print("command:", command)
+            ldict = {}
+            exec(command, globals(), ldict)
+
+            print('output:', ldict['output'])
+            print('type:', ldict['dtype'])
+            output = ldict['output']
+            dtype = ldict['dtype']
 
-example = 'Ex: How many rows are there?'
-user_input = st.text_input("Input your Query here", example)
+            if dtype == pd.core.series.Series and len(output) == 1:
+                command = f"output = {response.choices[0].text.replace('A: ', '')}.iloc[0]"
+                exec(command)
+                output = int(output)
 
-if user_input and user_input != example:
-    with st.spinner("Loading..."):
-        time.sleep(2)
+            if dtype in [int, np.int, np.int64, np.int32]:
+                st.text(output)
 
-        if user_input:
-            st.text(user_input)
+            if dtype in [pd.core.frame.DataFrame, pd.core.series.Series]:
+                st.dataframe(output)