Update app.py

app.py

@@ -1,40 +1,26 @@
 import streamlit as st
 import pandas as pd
-import json
-import os
+import sqlite3
 import plotly.express as px
-from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
-from datasets import Dataset
+import json
 
 # Set paths to the default files
-DEFAULT_PROMPT_PATH = "default_prompt.json"
+DEFAULT_PROMPT_PATH = "prompt_engineering.json"
 DEFAULT_METADATA_PATH = "default_metadata.csv"
 DEFAULT_DATA_PATH = "default_data.csv"
 
-# Load the RAG model with LLaMA-based retrieval-augmented generation
-@st.cache_resource(allow_output_mutation=True)
-def load_rag_model():
-    retriever = RagRetriever.from_pretrained("facebook/rag-token-base", index_name="custom")
-    tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-base")
-    model = RagSequenceForGeneration.from_pretrained("facebook/rag-token-base", retriever=retriever)
-    return model, tokenizer, retriever
-
-model, tokenizer, retriever = load_rag_model()
+# Load the prompt engineering JSON file (use default if no user-uploaded prompt file)
+with open(DEFAULT_PROMPT_PATH) as f:
+    prompt_data = json.load(f)
 
-# Title for the app
-st.title("Interactive Insights Chatbot with LLaMA + RAG")
-
-# Step 1: Upload prompt.json file (or use default)
-prompt_file = st.file_uploader("Upload your prompt.json file", type=["json"])
-if prompt_file is None:
-    with open(DEFAULT_PROMPT_PATH) as f:
-        prompt_data = json.load(f)
-        st.write("Using default prompt.json file.")
-else:
-    prompt_data = json.load(prompt_file)
-    st.write("Prompt JSON loaded successfully!")
+# Function to find a query based on the user prompt
+def get_query_from_prompt(user_prompt):
+    for item in prompt_data['prompts']:
+        if item['question'].lower() in user_prompt.lower():
+            return item['query']
+    return None  # Return None if no matching query is found
 
-# Step 2: Upload metadata.csv file (or use default)
+# Step 1: Upload metadata.csv file (or use default)
 metadata_file = st.file_uploader("Upload your metadata.csv file", type=["csv"])
 if metadata_file is None:
     metadata = pd.read_csv(DEFAULT_METADATA_PATH)
@@ -44,7 +30,7 @@ else:
     st.write("Metadata loaded successfully!")
     st.dataframe(metadata)
 
-# Step 3: Upload CSV data file (or use default)
+# Step 2: Upload CSV data file (or use default)
 csv_file = st.file_uploader("Upload your CSV file", type=["csv"])
 if csv_file is None:
     data = pd.read_csv(DEFAULT_DATA_PATH)
@@ -54,31 +40,24 @@ else:
     st.write("Data Preview:")
     st.dataframe(data.head())
 
-# Convert the CSV data to a Hugging Face Dataset for retrieval
-dataset = Dataset.from_pandas(data)
+# Step 3: Load CSV data into a SQLite database (SQL agent)
+conn = sqlite3.connect(':memory:')  # Use an in-memory SQLite database
+data.to_sql('sales_data', conn, index=False, if_exists='replace')
 
-# Step 4: Natural language prompt input
+# Step 4: Get user prompt and map to SQL query
 user_prompt = st.text_input("Enter your natural language prompt:")
 
-# Step 5: Process the prompt and generate insights using LLaMA + RAG
-if user_prompt and data is not None:
-    st.write(f"Processing your prompt: '{user_prompt}'")
-
-    # Tokenize the prompt for LLaMA + RAG
-    inputs = tokenizer(user_prompt, return_tensors="pt")
-
-    # Perform retrieval-augmented generation (RAG) by retrieving data from the dataset and generating the response
-    generated = model.generate(input_ids=inputs['input_ids'], num_return_sequences=1, num_beams=2)
-
-    # Decode the output from the LLaMA + RAG model
-    output = tokenizer.batch_decode(generated, skip_special_tokens=True)
-
-    st.write(f"Insights generated: {output[0]}")
-
-    # Example: if the prompt asks for a plot (like "show sales over time")
-    if "plot sales" in user_prompt.lower():
-        # Create an interactive bar chart
-        fig = px.bar(data, x='Date', y='Sales', title="Sales Over Time")
-        st.plotly_chart(fig)
+# Step 5: Process the prompt and generate SQL query dynamically
+if user_prompt:
+    query = get_query_from_prompt(user_prompt)
+    if query:
+        result = pd.read_sql(query, conn)
+        st.write("Query Results:")
+        st.dataframe(result)
+        
+        # If the query involves plotting (like "plot sales"), show the chart
+        if "plot" in user_prompt.lower():
+            fig = px.bar(result, x='Date', y='Sales', title="Sales Over Time")
+            st.plotly_chart(fig)
     else:
-        st.write("No recognized visual request in the prompt.")
+        st.write("Sorry, I couldn't find a matching query for your prompt.")