Update app.py

app.py

@@ -1,4 +1,3 @@
-import os
 import streamlit as st
 import pandas as pd
 from datasets import load_dataset
@@ -19,7 +18,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()
@@ -110,7 +109,7 @@ if "df" in st.session_state:
         df.to_csv("patent_data.csv", index=False)
 
         csv_agent = create_csv_agent(
-            ChatOpenAI(temperature=0, model="gpt-4", api_key=os.getenv("OPENAI_API_KEY")),
+            ChatOpenAI(temperature=0, model="gpt-4", api_key=st.secrets["api_key"]),
             path=["patent_data.csv"],
             verbose=True,
             agent_type=AgentType.OPENAI_FUNCTIONS,
@@ -123,9 +122,27 @@ if "df" in st.session_state:
     if st.button("Run Query"):
         with st.spinner("Running query..."):
             try:
-                result = csv_agent.invoke(query)
+                # Split query execution into smaller chunks if needed
+                max_rows = 1000
+                total_rows = len(df)
+                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=st.secrets["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("### Query Result:")
-                st.write(result)
+                st.write("\n".join(results))
+
             except Exception as e:
                 st.error(f"Error executing query: {e}")