Update app.py

app.py

@@ -19,7 +19,7 @@ def load_huggingface_dataset(dataset_name):
     try:
         # Incrementally update progress
         progress_bar.progress(10)
-        dataset = load_dataset(dataset_name, name="sample", split="train", trust_remote_code=True, uniform_split=True)
+        dataset = load_dataset(dataset_name, name="sample", split="train", trust_remote_code=True)
         progress_bar.progress(50)
         if hasattr(dataset, "to_pandas"):
             df = dataset.to_pandas()
@@ -123,9 +123,34 @@ if "df" in st.session_state:
     if st.button("Run Query"):
         with st.spinner("Running query..."):
             try:
-                result = csv_agent.invoke(query)
-                st.success("Query executed successfully!")
-                st.write("### Query Result:")
-                st.write(result)
+                # Check if the dataset is too large and split if needed
+                max_rows = 1000  # Limit chunk size to manage token limits
+                total_rows = len(df)
+
+                if total_rows > max_rows:
+                    results = []
+                    for start in range(0, total_rows, max_rows):
+                        chunk = df.iloc[start:start + max_rows]
+                        chunk.to_csv("chunk_data.csv", index=False)
+                        partial_agent = create_csv_agent(
+                            ChatOpenAI(temperature=0, model="gpt-4", api_key=os.getenv("OPENAI_API_KEY")),
+                            path=["chunk_data.csv"],
+                            verbose=True,
+                            agent_type=AgentType.OPENAI_FUNCTIONS,
+                            allow_dangerous_code=True
+                        )
+                        result = partial_agent.invoke(query)
+                        results.append(result)
+
+                    st.success("Query executed successfully!")
+                    st.write("### Combined Query Results:")
+                    st.write("\n".join(results))
+
+                else:
+                    result = csv_agent.invoke(query)
+                    st.success("Query executed successfully!")
+                    st.write("### Query Result:")
+                    st.write(result)
+
             except Exception as e:
-                st.error(f"Error executing query: {e}")
+                st.error(f"Error executing query: {e}")