Update app.py

app.py

@@ -1,69 +1,272 @@
-import streamlit as st 
-from lida import Manager, TextGenerationConfig , llm  
-from dotenv import load_dotenv
+import streamlit as st
+from lida import Manager, TextGenerationConfig, llm
+from lida.datamodel import Goal
 import os
-import openai
-from PIL import Image
-from io import BytesIO
-import base64
-
-load_dotenv()
-openai.api_key = os.getenv('OPENAI_API_KEY')
-
-def base64_to_image(base64_string):
-    # Decode the base64 string
-    byte_data = base64.b64decode(base64_string)
-    
-    # Use BytesIO to convert the byte data to image
-    return Image.open(BytesIO(byte_data))
-
-
-lida = Manager(text_gen = llm("openai"))
-textgen_config = TextGenerationConfig(n=1, temperature=0.5, model="gpt-3.5-turbo-0301", use_cache=True)
-
-menu = st.sidebar.selectbox("Choose an Option", ["Summarize", "Question based Graph"])
-
-if menu == "Summarize":
-    st.subheader("Summarization of your Data")
-    file_uploader = st.file_uploader("Upload your CSV", type="csv")
-    if file_uploader is not None:
-        path_to_save = "filename.csv"
-        with open(path_to_save, "wb") as f:
-            f.write(file_uploader.getvalue())
-        summary = lida.summarize("filename.csv", summary_method="default", textgen_config=textgen_config)
-        st.write(summary)
-        goals = lida.goals(summary, n=2, textgen_config=textgen_config)
-        for goal in goals:
-            st.write(goal)
-        i = 0
-        library = "seaborn"
-        textgen_config = TextGenerationConfig(n=1, temperature=0.2, use_cache=True)
-        charts = lida.visualize(summary=summary, goal=goals[i], textgen_config=textgen_config, library=library)  
-        img_base64_string = charts[0].raster
-        img = base64_to_image(img_base64_string)
-        st.image(img)
-        
-
-
-        
-elif menu == "Question based Graph":
-    st.subheader("Query your Data to Generate Graph")
-    file_uploader = st.file_uploader("Upload your CSV", type="csv")
-    if file_uploader is not None:
-        path_to_save = "filename1.csv"
-        with open(path_to_save, "wb") as f:
-            f.write(file_uploader.getvalue())
-        text_area = st.text_area("Query your Data to Generate Graph", height=200)
-        if st.button("Generate Graph"):
-            if len(text_area) > 0:
-                st.info("Your Query: " + text_area)
-                lida = Manager(text_gen = llm("openai")) 
-                textgen_config = TextGenerationConfig(n=1, temperature=0.2, use_cache=True)
-                summary = lida.summarize("filename1.csv", summary_method="default", textgen_config=textgen_config)
-                user_query = text_area
-                charts = lida.visualize(summary=summary, goal=user_query, textgen_config=textgen_config)  
-                charts[0]
-                image_base64 = charts[0].raster
-                img = base64_to_image(image_base64)
-                st.image(img)
-            
+import pandas as pd
+
+
+# make data dir if it doesn't exist
+os.makedirs("data", exist_ok=True)
+
+st.set_page_config(
+    page_title="LIDA: Automatic Generation of Visualizations and Infographics",
+    page_icon="📊",
+)
+
+st.write("# LIDA: Automatic Generation of Visualizations and Infographics using Large Language Models 📊")
+
+st.sidebar.write("## Setup")
+
+# Step 1 - Get OpenAI API key
+openai_key = os.getenv("OPENAI_API_KEY")
+
+if not openai_key:
+    openai_key = st.sidebar.text_input("Enter OpenAI API key:")
+    if openai_key:
+        display_key = openai_key[:2] + "*" * (len(openai_key) - 5) + openai_key[-3:]
+        st.sidebar.write(f"Current key: {display_key}")
+    else:
+        st.sidebar.write("Please enter OpenAI API key.")
+else:
+    display_key = openai_key[:2] + "*" * (len(openai_key) - 5) + openai_key[-3:]
+    st.sidebar.write(f"OpenAI API key loaded from environment variable: {display_key}")
+
+st.markdown(
+    """
+    LIDA is a library for generating data visualizations and data-faithful infographics.
+    LIDA is grammar agnostic (will work with any programming language and visualization
+    libraries e.g. matplotlib, seaborn, altair, d3 etc) and works with multiple large language
+    model providers (OpenAI, Azure OpenAI, PaLM, Cohere, Huggingface). Details on the components
+    of LIDA are described in the [paper here](https://arxiv.org/abs/2303.02927) and in this
+    tutorial [notebook](notebooks/tutorial.ipynb). See the project page [here](https://microsoft.github.io/lida/) for updates!.
+
+   This demo shows how to use the LIDA python api with Streamlit. [More](/about).
+
+   ----
+""")
+
+# Step 2 - Select a dataset and summarization method
+if openai_key:
+    # Initialize selected_dataset to None
+    selected_dataset = None
+
+    # select model from gpt-4 , gpt-3.5-turbo, gpt-3.5-turbo-16k
+    st.sidebar.write("## Text Generation Model")
+    models = ["gpt-4", "gpt-3.5-turbo", "gpt-3.5-turbo-16k"]
+    selected_model = st.sidebar.selectbox(
+        'Choose a model',
+        options=models,
+        index=0
+    )
+
+    # select temperature on a scale of 0.0 to 1.0
+    # st.sidebar.write("## Text Generation Temperature")
+    temperature = st.sidebar.slider(
+        "Temperature",
+        min_value=0.0,
+        max_value=1.0,
+        value=0.0)
+
+    # set use_cache in sidebar
+    use_cache = st.sidebar.checkbox("Use cache", value=True)
+
+    # Handle dataset selection and upload
+    st.sidebar.write("## Data Summarization")
+    st.sidebar.write("### Choose a dataset")
+
+    datasets = [
+        {"label": "Select a dataset", "url": None},
+        {"label": "Cars", "url": "https://raw.githubusercontent.com/uwdata/draco/master/data/cars.csv"},
+        {"label": "Weather", "url": "https://raw.githubusercontent.com/uwdata/draco/master/data/weather.json"},
+    ]
+
+    selected_dataset_label = st.sidebar.selectbox(
+        'Choose a dataset',
+        options=[dataset["label"] for dataset in datasets],
+        index=0
+    )
+
+    upload_own_data = st.sidebar.checkbox("Upload your own data")
+
+    if upload_own_data:
+        uploaded_file = st.sidebar.file_uploader("Choose a CSV or JSON file", type=["csv", "json"])
+
+        if uploaded_file is not None:
+            # Get the original file name and extension
+            file_name, file_extension = os.path.splitext(uploaded_file.name)
+
+            # Load the data depending on the file type
+            if file_extension.lower() == ".csv":
+                data = pd.read_csv(uploaded_file)
+            elif file_extension.lower() == ".json":
+                data = pd.read_json(uploaded_file)
+
+            # Save the data using the original file name in the data dir
+            uploaded_file_path = os.path.join("data", uploaded_file.name)
+            data.to_csv(uploaded_file_path, index=False)
+
+            selected_dataset = uploaded_file_path
+
+            datasets.append({"label": file_name, "url": uploaded_file_path})
+
+            # st.sidebar.write("Uploaded file path: ", uploaded_file_path)
+    else:
+        selected_dataset = datasets[[dataset["label"]
+                                     for dataset in datasets].index(selected_dataset_label)]["url"]
+
+    if not selected_dataset:
+        st.info("To continue, select a dataset from the sidebar on the left or upload your own.")
+
+    st.sidebar.write("### Choose a summarization method")
+    # summarization_methods = ["default", "llm", "columns"]
+    summarization_methods = [
+        {"label": "llm",
+         "description":
+         "Uses the LLM to generate annotate the default summary, adding details such as semantic types for columns and dataset description"},
+        {"label": "default",
+         "description": "Uses dataset column statistics and column names as the summary"},
+
+        {"label": "columns", "description": "Uses the dataset column names as the summary"}]
+
+    # selected_method = st.sidebar.selectbox("Choose a method", options=summarization_methods)
+    selected_method_label = st.sidebar.selectbox(
+        'Choose a method',
+        options=[method["label"] for method in summarization_methods],
+        index=0
+    )
+
+    selected_method = summarization_methods[[
+        method["label"] for method in summarization_methods].index(selected_method_label)]["label"]
+
+    # add description of selected method in very small font to sidebar
+    selected_summary_method_description = summarization_methods[[
+        method["label"] for method in summarization_methods].index(selected_method_label)]["description"]
+
+    if selected_method:
+        st.sidebar.markdown(
+            f"<span> {selected_summary_method_description} </span>",
+            unsafe_allow_html=True)
+
+# Step 3 - Generate data summary
+if openai_key and selected_dataset and selected_method:
+    lida = Manager(text_gen=llm("openai", api_key=openai_key))
+    textgen_config = TextGenerationConfig(
+        n=1,
+        temperature=temperature,
+        model=selected_model,
+        use_cache=use_cache)
+
+    st.write("## Summary")
+    # **** lida.summarize *****
+    summary = lida.summarize(
+        selected_dataset,
+        summary_method=selected_method,
+        textgen_config=textgen_config)
+
+    if "dataset_description" in summary:
+        st.write(summary["dataset_description"])
+
+    if "fields" in summary:
+        fields = summary["fields"]
+        nfields = []
+        for field in fields:
+            flatted_fields = {}
+            flatted_fields["column"] = field["column"]
+            # flatted_fields["dtype"] = field["dtype"]
+            for row in field["properties"].keys():
+                if row != "samples":
+                    flatted_fields[row] = field["properties"][row]
+                else:
+                    flatted_fields[row] = str(field["properties"][row])
+            # flatted_fields = {**flatted_fields, **field["properties"]}
+            nfields.append(flatted_fields)
+        nfields_df = pd.DataFrame(nfields)
+        st.write(nfields_df)
+    else:
+        st.write(str(summary))
+
+    # Step 4 - Generate goals
+    if summary:
+        st.sidebar.write("### Goal Selection")
+
+        num_goals = st.sidebar.slider(
+            "Number of goals to generate",
+            min_value=1,
+            max_value=10,
+            value=4)
+        own_goal = st.sidebar.checkbox("Add Your Own Goal")
+
+        # **** lida.goals *****
+        goals = lida.goals(summary, n=num_goals, textgen_config=textgen_config)
+        st.write(f"## Goals ({len(goals)})")
+
+        default_goal = goals[0].question
+        goal_questions = [goal.question for goal in goals]
+
+        if own_goal:
+            user_goal = st.sidebar.text_input("Describe Your Goal")
+
+            if user_goal:
+
+                new_goal = Goal(question=user_goal, visualization=user_goal, rationale="")
+                goals.append(new_goal)
+                goal_questions.append(new_goal.question)
+
+        selected_goal = st.selectbox('Choose a generated goal', options=goal_questions, index=0)
+
+        # st.markdown("### Selected Goal")
+        selected_goal_index = goal_questions.index(selected_goal)
+        st.write(goals[selected_goal_index])
+
+        selected_goal_object = goals[selected_goal_index]
+
+        # Step 5 - Generate visualizations
+        if selected_goal_object:
+            st.sidebar.write("## Visualization Library")
+            visualization_libraries = ["seaborn", "matplotlib", "plotly"]
+
+            selected_library = st.sidebar.selectbox(
+                'Choose a visualization library',
+                options=visualization_libraries,
+                index=0
+            )
+
+            # Update the visualization generation call to use the selected library.
+            st.write("## Visualizations")
+
+            # slider for number of visualizations
+            num_visualizations = st.sidebar.slider(
+                "Number of visualizations to generate",
+                min_value=1,
+                max_value=10,
+                value=2)
+
+            textgen_config = TextGenerationConfig(
+                n=num_visualizations, temperature=temperature,
+                model=selected_model,
+                use_cache=use_cache)
+
+            # **** lida.visualize *****
+            visualizations = lida.visualize(
+                summary=summary,
+                goal=selected_goal_object,
+                textgen_config=textgen_config,
+                library=selected_library)
+
+            viz_titles = [f'Visualization {i+1}' for i in range(len(visualizations))]
+
+            selected_viz_title = st.selectbox('Choose a visualization', options=viz_titles, index=0)
+
+            selected_viz = visualizations[viz_titles.index(selected_viz_title)]
+
+            if selected_viz.raster:
+                from PIL import Image
+                import io
+                import base64
+
+                imgdata = base64.b64decode(selected_viz.raster)
+                img = Image.open(io.BytesIO(imgdata))
+                st.image(img, caption=selected_viz_title, use_column_width=True)
+
+            st.write("### Visualization Code")
+            st.code(selected_viz.code)