Update app.py

app.py

@@ -7,7 +7,7 @@ from io import StringIO
 from transformers import pipeline
 
 # Load a lightweight NLP model for query understanding
-nlp = pipeline("text-classification", model="distilbert-base-uncased", tokenizer="distilbert-base-uncased")
+nlp = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
 
 # Function to load the uploaded file (CSV or Excel)
 def load_file(uploaded_file):
@@ -34,11 +34,11 @@ def infer_column(data, synonyms):
     return None
 
 # Function to classify the user query
-def classify_query(query):
+def classify_query(query, candidate_labels):
     """Classify the user query into graph types."""
-    results = nlp(query)
+    results = nlp(query, candidate_labels)
     if results:
-        return results[0]['label']
+        return results['labels'][0]
     return None
 
 # Function to generate graph based on user query
@@ -47,55 +47,52 @@ def generate_graph(data, query):
     try:
         fig, ax = plt.subplots(figsize=(10, 6))
 
-        # Infer column names
-        country_col = infer_column(data, {"country", "countries"})
-        sales_col = infer_column(data, {"gross_sales", "sales", "revenue"})
-        date_col = infer_column(data, {"date", "time"})
+        # Infer column names dynamically
+        numerical_columns = data.select_dtypes(include=['number']).columns.tolist()
+        categorical_columns = data.select_dtypes(include=['object', 'category']).columns.tolist()
+        datetime_columns = data.select_dtypes(include=['datetime']).columns.tolist()
+
+        # Define possible graph types
+        candidate_labels = ["bar chart", "line chart", "scatter plot", "histogram"]
+        query_type = classify_query(query, candidate_labels)
+
+        if query_type == "bar chart" and categorical_columns and numerical_columns:
+            # Bar chart for categorical vs numerical
+            x_col = st.selectbox("Select the categorical column:", categorical_columns)
+            y_col = st.selectbox("Select the numerical column:", numerical_columns)
+            aggregated_data = data[[x_col, y_col]].groupby(x_col).sum().reset_index()
+            sns.barplot(x=x_col, y=y_col, data=aggregated_data, ax=ax, color='skyblue')
+            ax.set_xticklabels(ax.get_xticklabels(), rotation=45, ha='right')
+            ax.set_title(f"Bar Chart: {x_col} vs {y_col}")
+            st.pyplot(fig)
 
-        query_type = classify_query(query)
+        elif query_type == "line chart" and datetime_columns and numerical_columns:
+            # Line chart for numerical trend over time
+            x_col = st.selectbox("Select the datetime column:", datetime_columns)
+            y_col = st.selectbox("Select the numerical column:", numerical_columns)
+            data[x_col] = pd.to_datetime(data[x_col])
+            trend_data = data.groupby(x_col)[y_col].sum().reset_index()
+            sns.lineplot(x=x_col, y=y_col, data=trend_data, ax=ax)
+            ax.set_title(f"Line Chart: {y_col} Over Time")
+            st.pyplot(fig)
 
-        if "bar" in query.lower() and country_col and sales_col:
-            # Bar chart for countries and gross sales
-            country_data = data[[country_col, sales_col]].groupby(country_col).sum().reset_index()
-            sns.barplot(x=country_col, y=sales_col, data=country_data, ax=ax, color='skyblue')
-            ax.set_xticklabels(ax.get_xticklabels(), rotation=45, ha='right')
-            ax.set_title(f"Bar Chart: {country_col} vs {sales_col}")
+        elif query_type == "scatter plot" and len(numerical_columns) >= 2:
+            # Scatter plot for numerical relationships
+            x_col = st.selectbox("Select the x-axis numerical column:", numerical_columns)
+            y_col = st.selectbox("Select the y-axis numerical column:", numerical_columns)
+            sns.scatterplot(x=x_col, y=y_col, data=data, ax=ax)
+            ax.set_title(f"Scatter Plot: {x_col} vs {y_col}")
             st.pyplot(fig)
-        elif "line" in query.lower() and date_col and sales_col:
-            # Line chart for sales trend over time
-            data[date_col] = pd.to_datetime(data[date_col])
-            sales_trend = data.groupby(date_col)[sales_col].sum().reset_index()
-            sns.lineplot(x=date_col, y=sales_col, data=sales_trend, ax=ax)
-            ax.set_title(f"Line Chart: {sales_col} Over Time")
+
+        elif query_type == "histogram" and numerical_columns:
+            # Histogram for a numerical column
+            hist_col = st.selectbox("Select the numerical column:", numerical_columns)
+            sns.histplot(data[hist_col], bins=20, kde=True, ax=ax, color='green')
+            ax.set_title(f"Histogram of {hist_col}")
             st.pyplot(fig)
-        elif "scatter" in query.lower():
-            # Scatter plot for relationships
-            if "between" in query.lower():
-                columns = query.lower().split("between")[-1].strip().split("and")
-                if len(columns) == 2:
-                    x_col = infer_column(data, {columns[0].strip()})
-                    y_col = infer_column(data, {columns[1].strip()})
-
-                    if x_col and y_col:
-                        sns.scatterplot(x=x_col, y=y_col, data=data, ax=ax)
-                        ax.set_title(f"Scatter Plot: {x_col} vs {y_col}")
-                        st.pyplot(fig)
-                        return
-            st.error("Please specify valid columns for the scatter plot.")
-        elif "histogram" in query.lower():
-            # Histogram for a specified column
-            if "for" in query.lower():
-                column = query.lower().split("for")[-1].strip()
-                hist_col = infer_column(data, {column})
-
-                if hist_col:
-                    sns.histplot(data[hist_col], bins=20, kde=True, ax=ax, color='green')
-                    ax.set_title(f"Histogram of {hist_col}")
-                    st.pyplot(fig)
-                    return
-            st.error("Please specify a valid column for the histogram.")
+
         else:
-            st.error("Unsupported graph type. Try asking for a bar chart, line chart, scatter plot, or histogram.")
+            st.error("Unsupported graph type or insufficient data. Try asking for a bar chart, line chart, scatter plot, or histogram.")
     except Exception as e:
         st.error(f"Error generating graph: {e}")