Update app.py

app.py

@@ -7,42 +7,46 @@ from langchain_community.vectorstores import FAISS
 from langchain_openai import ChatOpenAI
 from langchain.chains import RetrievalQA
 from langchain.schema import Document
-from datasets import load_dataset  
+from datasets import load_dataset
 import os
 
+
+# Function to load dataset into session
+def load_dataset_into_session():
+    input_option = st.sidebar.radio("Select Dataset Input:", ["Use Hugging Face Dataset", "Upload CSV File"])
+    if input_option == "Use Hugging Face Dataset":
+        dataset_name = st.sidebar.text_input("Enter Hugging Face Dataset Name:", value="HUPD/hupd")
+        if st.sidebar.button("Load Dataset"):
+            try:
+                # Ensure `uniform_split=True` is provided for compatibility
+                dataset = load_dataset(dataset_name, name="sample", split="train", trust_remote_code=True, uniform_split=True)
+                st.session_state.df = pd.DataFrame(dataset)
+                st.sidebar.success(f"Dataset '{dataset_name}' loaded successfully!")
+            except Exception as e:
+                st.sidebar.error(f"Error loading dataset: {e}")
+    elif input_option == "Upload CSV File":
+        uploaded_file = st.sidebar.file_uploader("Upload CSV File:", type=["csv"])
+        if uploaded_file:
+            try:
+                st.session_state.df = pd.read_csv(uploaded_file)
+                st.sidebar.success("File uploaded successfully!")
+            except Exception as e:
+                st.sidebar.error(f"Error loading file: {e}")
+
+
 # Title
-st.title("PandasAI Data Analysis Tool with RAG")
+st.title("XYZ")
 
 # Fetch API keys from environment variables
 api_key = os.getenv("OPENAI_API_KEY")
 pandasai_api_key = os.getenv("PANDASAI_API_KEY")
 
-# Dataset selection
-st.sidebar.header("Dataset Input Options")
-input_option = st.sidebar.radio("Select Dataset Input:", ["Use Hugging Face Dataset", "Upload CSV File"])
-
 # Initialize session state for the dataframe
 if "df" not in st.session_state:
     st.session_state.df = None
 
-# Dataset loading logic
-if input_option == "Use Hugging Face Dataset":
-    dataset_name = st.sidebar.text_input("Enter Hugging Face Dataset Name:", value="HUPD/hupd")
-    if st.sidebar.button("Load Dataset"):
-        try:
-            dataset = load_dataset(dataset_name, name="sample", split="train", trust_remote_code=True)
-            st.session_state.df = pd.DataFrame(dataset)
-            st.sidebar.success(f"Dataset '{dataset_name}' loaded successfully!")
-        except Exception as e:
-            st.sidebar.error(f"Error loading dataset: {e}")
-elif input_option == "Upload CSV File":
-    uploaded_file = st.sidebar.file_uploader("Upload CSV File:", type=["csv"])
-    if uploaded_file:
-        try:
-            st.session_state.df = pd.read_csv(uploaded_file)
-            st.sidebar.success("File uploaded successfully!")
-        except Exception as e:
-            st.sidebar.error(f"Error loading file: {e}")
+# Call the function to load the dataset
+load_dataset_into_session()
 
 # Show the loaded dataframe preview
 if st.session_state.df is not None:
@@ -121,4 +125,4 @@ else:
 if not api_key:
     st.error("Missing OpenAI API Key in environment variables.")
 if not pandasai_api_key:
-    st.error("Missing PandasAI API Key in environment variables.")
+    st.error("Missing PandasAI API Key in environment variables.")