Update app.py

app.py

@@ -1,4 +1,5 @@
 import os
+import tempfile
 import streamlit as st
 import pandas as pd
 from datasets import load_dataset
@@ -19,7 +20,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()
@@ -107,15 +108,16 @@ if "df" in st.session_state:
     st.header("Run Queries on Patent Data")
 
     with st.spinner("Setting up LangChain CSV Agent..."):
-        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")),
-            path=["patent_data.csv"],
-            verbose=True,
-            agent_type=AgentType.OPENAI_FUNCTIONS,
-            allow_dangerous_code=True
-        )
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".csv") as temp_file:
+            df.to_csv(temp_file.name, index=False)
+
+            csv_agent = create_csv_agent(
+                ChatOpenAI(temperature=0, model="gpt-4", api_key=os.getenv("OPENAI_API_KEY")),
+                path=[temp_file.name],
+                verbose=True,
+                agent_type=AgentType.OPENAI_FUNCTIONS,
+                allow_dangerous_code=True
+            )
 
     # Query Input and Execution
     query = st.text_area("Enter your natural language query:", "How many patents are related to AI?")
@@ -123,24 +125,27 @@ if "df" in st.session_state:
     if st.button("Run Query"):
         with st.spinner("Running query..."):
             try:
-                # Check if the dataset is too large and split if needed
-                max_rows = 1000  # Limit chunk size to manage token limits
+                # Token limit configuration
+                max_rows = 1000  # Adjust chunk size dynamically
                 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)
+                        with tempfile.NamedTemporaryFile(delete=False, suffix=".csv") as chunk_file:
+                            chunk.to_csv(chunk_file.name, index=False)
+
+                            # Update the agent dynamically with the chunk
+                            csv_agent = create_csv_agent(
+                                ChatOpenAI(temperature=0, model="gpt-4", api_key=os.getenv("OPENAI_API_KEY")),
+                                path=[chunk_file.name],
+                                verbose=False,
+                                agent_type=AgentType.OPENAI_FUNCTIONS,
+                                allow_dangerous_code=False
+                            )
+                            result = csv_agent.invoke(query)
+                            results.append(result)
 
                     st.success("Query executed successfully!")
                     st.write("### Combined Query Results:")