Spaces:
Sleeping
Sleeping
File size: 3,846 Bytes
b193f65 |
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 |
import streamlit as st
from pandasai.llm.openai import OpenAI
from dotenv import load_dotenv
import os
import pandas as pd
from pandasai import PandasAI
from datasets import load_dataset
import time
openai_api_key = os.getenv("OPENAI_API_KEY")
def chat_with_csv(df, prompt):
llm = OpenAI(api_token=openai_api_key)
pandas_ai = PandasAI(llm)
result = pandas_ai.run(df, prompt=prompt)
return result
def load_huggingface_dataset(dataset_name):
progress_bar = st.progress(0)
try:
progress_bar.progress(10)
dataset = load_dataset(dataset_name, name="sample", split="train", trust_remote_code=True, uniform_split=True)
progress_bar.progress(50)
if hasattr(dataset, "to_pandas"):
df = dataset.to_pandas()
else:
df = pd.DataFrame(dataset)
progress_bar.progress(100)
return df
except Exception as e:
progress_bar.progress(0)
raise e
def load_uploaded_csv(uploaded_file):
progress_bar = st.progress(0)
try:
progress_bar.progress(10)
time.sleep(1)
progress_bar.progress(50)
df = pd.read_csv(uploaded_file)
progress_bar.progress(100)
return df
except Exception as e:
progress_bar.progress(0)
raise e
def load_dataset_into_session():
input_option = st.radio(
"Select Dataset Input:",
["Use Repo Directory Dataset", "Use Hugging Face Dataset", "Upload CSV File"],
index=1,
horizontal=True
)
if input_option == "Use Repo Directory Dataset":
file_path = "./source/test.csv"
if st.button("Load Dataset"):
try:
with st.spinner("Loading dataset from the repo directory..."):
st.session_state.df = pd.read_csv(file_path)
st.success(f"File loaded successfully from '{file_path}'!")
except Exception as e:
st.error(f"Error loading dataset from the repo directory: {e}")
elif input_option == "Use Hugging Face Dataset":
dataset_name = st.text_input("Enter Hugging Face Dataset Name:", value="HUPD/hupd")
if st.button("Load Dataset"):
try:
st.session_state.df = load_huggingface_dataset(dataset_name)
st.success(f"Hugging Face Dataset '{dataset_name}' loaded successfully!")
except Exception as e:
st.error(f"Error loading Hugging Face dataset: {e}")
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 = load_uploaded_csv(uploaded_file)
st.success("File uploaded successfully!")
except Exception as e:
st.error(f"Error reading uploaded file: {e}")
# Streamlit app main
st.set_page_config(layout='wide')
st.title("ChatCSV powered by LLM")
# Ensure session state for the dataframe
if "df" not in st.session_state:
st.session_state.df = pd.DataFrame() # Initialize with an empty dataframe
st.header("Load Your Dataset")
load_dataset_into_session()
if not st.session_state.df.empty:
st.subheader("Dataset Preview")
st.dataframe(st.session_state.df, use_container_width=True)
st.subheader("Chat with Your Dataset")
user_query = st.text_area("Enter your query:")
if st.button("Run Query"):
if user_query.strip():
with st.spinner("Processing your query..."):
try:
result = chat_with_csv(st.session_state.df, user_query)
st.success(result)
except Exception as e:
st.error(f"Error processing your query: {e}")
else:
st.warning("Please enter a query before running.")
|