File size: 7,799 Bytes
39a3df1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
import streamlit as st
import pandas as pd
import plotly.express as px
from pandasai import Agent
from langchain_community.embeddings.openai import OpenAIEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_openai import ChatOpenAI
from langchain.chains import RetrievalQA
from langchain.schema import Document
import os
import re

# Set title
st.title("Data Analyzer")

# API keys
api_key = os.getenv("OPENAI_API_KEY")
pandasai_api_key = os.getenv("PANDASAI_API_KEY")

if not api_key or not pandasai_api_key:
    st.warning("API keys for OpenAI or PandasAI are missing. Ensure both keys are set in environment variables.")

# Add session reset button
if st.button("Reset Session"):
    for key in list(st.session_state.keys()):
        del st.session_state[key]
    st.experimental_rerun()

# Function to validate and clean dataset
def validate_and_clean_dataset(df):
    # Rename columns for consistency
    df.columns = [col.strip().lower().replace(" ", "_") for col in df.columns]
    # Check for missing values
    if df.isnull().values.any():
        st.warning("Dataset contains missing values. Consider cleaning the data.")
    return df

# Function to load datasets into session
def load_dataset_into_session():
    input_option = st.radio(
        "Select Dataset Input:",
        ["Use Repo Directory Dataset", "Use Hugging Face Dataset", "Upload CSV File"],
    )

    # Option 1: Load dataset from the repo directory
    if input_option == "Use Repo Directory Dataset":
        file_path = "./source/test.csv"
        if st.button("Load Dataset"):
            try:
                st.session_state.df = pd.read_csv(file_path)
                st.session_state.df = validate_and_clean_dataset(st.session_state.df)
                st.success(f"File loaded successfully from '{file_path}'!")
            except Exception as e:
                st.error(f"Error loading dataset from the repo directory: {e}")

    # Option 2: Load dataset from Hugging Face
    elif input_option == "Use Hugging Face Dataset":
        dataset_name = st.text_input(
            "Enter Hugging Face Dataset Name:", value="HUPD/hupd"
        )
        if st.button("Load Hugging Face Dataset"):
            try:
                from datasets import load_dataset
                dataset = load_dataset(dataset_name, split="train", trust_remote_code=True)
                if hasattr(dataset, "to_pandas"):
                    st.session_state.df = dataset.to_pandas()
                else:
                    st.session_state.df = pd.DataFrame(dataset)
                st.session_state.df = validate_and_clean_dataset(st.session_state.df)
                st.success(f"Hugging Face Dataset '{dataset_name}' loaded successfully!")
            except Exception as e:
                st.error(f"Error loading Hugging Face dataset: {e}")

    # Option 3: Upload CSV File
    elif input_option == "Upload CSV File":
        uploaded_file = st.file_uploader("Upload a CSV File:", type=["csv"])
        if uploaded_file:
            try:
                st.session_state.df = pd.read_csv(uploaded_file)
                st.session_state.df = validate_and_clean_dataset(st.session_state.df)
                st.success("File uploaded successfully!")
            except Exception as e:
                st.error(f"Error reading uploaded file: {e}")

load_dataset_into_session()

# Check if the dataset and API keys are loaded
if "df" in st.session_state and api_key and pandasai_api_key:
    # Set API keys
    os.environ["OPENAI_API_KEY"] = api_key
    os.environ["PANDASAI_API_KEY"] = pandasai_api_key

    df = st.session_state.df
    st.write("Dataset Preview:")
    st.write(df.head())  # Ensure the dataset preview is displayed only once

    # Set up PandasAI Agent
    try:
        agent = Agent(df)
        st.info("PandasAI Agent initialized successfully.")
    except Exception as e:
        st.error(f"Error initializing PandasAI Agent: {str(e)}")

    # Convert dataframe into documents
    try:
        documents = [
            Document(
                page_content=", ".join([f"{col}: {row[col]}" for col in df.columns]),
                metadata={"index": index}
            )
            for index, row in df.iterrows()
        ]
        st.info("Documents created successfully for RAG.")
    except Exception as e:
        st.error(f"Error creating documents for RAG: {str(e)}")

    # Set up RAG
    try:
        embeddings = OpenAIEmbeddings()
        vectorstore = FAISS.from_documents(documents, embeddings)
        retriever = vectorstore.as_retriever()
        qa_chain = RetrievalQA.from_chain_type(
            llm=ChatOpenAI(),
            chain_type="stuff",
            retriever=retriever
        )
        st.info("RAG setup completed successfully.")
    except Exception as e:
        st.error(f"Error setting up RAG: {str(e)}")

    # Create tabs
    tab1, tab2, tab3 = st.tabs(["PandasAI Analysis", "RAG Q&A", "Data Visualization"])

    with tab1:
        st.subheader("Data Analysis with PandasAI")
        pandas_question = st.text_input("Ask a question about the dataset (PandasAI):")
        if pandas_question:
            try:
                result = agent.chat(pandas_question)
                st.write("PandasAI Answer:", result)
                if hasattr(agent, "last_output"):
                    st.write("PandasAI Intermediate Output:", agent.last_output)
            except Exception as e:
                st.error(f"PandasAI encountered an error: {str(e)}")
                # Fallback: Direct pandas filtering
                if "patent_number" in pandas_question.lower() and "decision" in pandas_question.lower():
                    try:
                        match = re.search(r'\d{7,}', pandas_question)
                        if match:
                            patent_number = match.group()
                            decision = df.loc[df['patent_number'] == int(patent_number), 'decision']
                            st.write(f"Fallback Answer: The decision for patent {patent_number} is '{decision.iloc[0]}'.")
                        else:
                            st.write("Could not extract patent number from the query.")
                    except Exception as fallback_error:
                        st.error(f"Fallback processing failed: {fallback_error}")

    with tab2:
        st.subheader("Q&A with RAG")
        rag_question = st.text_input("Ask a question about the dataset (RAG):")
        if rag_question:
            try:
                result = qa_chain.run(rag_question)
                st.write("RAG Answer:", result)
            except Exception as e:
                st.error(f"RAG encountered an error: {str(e)}")

    with tab3:
        st.subheader("Data Visualization")
        viz_question = st.text_input("What kind of graph would you like? (e.g., 'Show a scatter plot of salary vs experience')")
        if viz_question:
            try:
                result = agent.chat(viz_question)
                code_pattern = r'```python\n(.*?)\n```'
                code_match = re.search(code_pattern, result, re.DOTALL)
                
                if code_match:
                    viz_code = code_match.group(1)
                    exec(viz_code)
                else:
                    st.write("Unable to generate the graph. Showing fallback example.")
                    fig = px.scatter(df, x=df.columns[0], y=df.columns[1])
                    st.plotly_chart(fig)
            except Exception as e:
                st.error(f"An error occurred during visualization: {str(e)}")
else:
    if not api_key:
        st.warning("Please set the OpenAI API key in environment variables.")
    if not pandasai_api_key:
        st.warning("Please set the PandasAI API key in environment variables.")