Spaces:
Running
Running
| import streamlit as st | |
| from pandasai.llm.openai import OpenAI | |
| from pandasai import PandasAI | |
| from dotenv import load_dotenv | |
| import os | |
| import pandas as pd | |
| import time | |
| from datasets import load_dataset # For Hugging Face datasets | |
| # Load environment variables | |
| load_dotenv() | |
| openai_api_key = os.getenv("OPENAI_API_KEY") | |
| groq_api_key = os.getenv("GROQ_API_KEY") | |
| # Initialize the LLM based on user selection | |
| def initialize_llm(model_choice): | |
| if model_choice == "llama-3.3-70b": | |
| if not groq_api_key: | |
| st.error("Groq API key is missing. Please set the GROQ_API_KEY environment variable.") | |
| return None | |
| return ChatGroq(groq_api_key=groq_api_key, model="groq/llama-3.3-70b-versatile") | |
| elif model_choice == "GPT-4o": | |
| if not openai_api_key: | |
| st.error("OpenAI API key is missing. Please set the OPENAI_API_KEY environment variable.") | |
| return None | |
| return OpenAI(api_token=openai_api_key) | |
| # Load Hugging Face dataset | |
| 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 | |
| # Load uploaded CSV file | |
| 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 | |
| # Main function to handle interactions | |
| def chat_with_csv(df, prompt, llm): | |
| pandas_ai = PandasAI(llm) | |
| result = pandas_ai.run(df, prompt=prompt) | |
| return result | |
| # Streamlit app layout | |
| st.set_page_config(layout="wide") | |
| st.title("ChatCSV powered by LLM") | |
| # Initialize session state for data storage | |
| if "data" not in st.session_state: | |
| st.session_state.data = None | |
| if "llm" not in st.session_state: | |
| st.session_state.llm = None | |
| # Select LLM model | |
| model_choice = st.radio("Select LLM", ["GPT-4o", "llama-3.3-70b"], index=0, horizontal=True) | |
| st.session_state.llm = initialize_llm(model_choice) | |
| if st.session_state.llm is not None: | |
| # Dataset selection | |
| input_option = st.radio( | |
| "Select Dataset Input:", | |
| ["Use Repo Directory Dataset", "Use Hugging Face Dataset", "Upload CSV File"], | |
| index=2, | |
| 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.data = 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.data = 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": | |
| input_csv = st.file_uploader("Upload a CSV File:", type=["csv"]) | |
| if input_csv: | |
| try: | |
| st.session_state.data = load_uploaded_csv(input_csv) | |
| st.success("File uploaded successfully!") | |
| except Exception as e: | |
| st.error(f"Error reading uploaded file: {e}") | |
| # Display the loaded dataset | |
| if st.session_state.data is not None: | |
| st.dataframe(st.session_state.data, use_container_width=True) | |
| # Chat interface | |
| input_text = st.text_area("Enter your query") | |
| if input_text and st.button("Chat with CSV"): | |
| try: | |
| st.info("Your Query: " + input_text) | |
| result = chat_with_csv(st.session_state.data, input_text, st.session_state.llm) | |
| st.success(result) | |
| except Exception as e: | |
| st.error(f"Error processing your query: {e}") | |