Update app.py

app.py

@@ -4,6 +4,10 @@ import matplotlib.pyplot as plt
 import seaborn as sns
 import os
 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")
 
 # Function to load the uploaded file (CSV or Excel)
 def load_file(uploaded_file):
@@ -29,6 +33,14 @@ def infer_column(data, synonyms):
             return column
     return None
 
+# Function to classify the user query
+def classify_query(query):
+    """Classify the user query into graph types."""
+    results = nlp(query)
+    if results:
+        return results[0]['label']
+    return None
+
 # Function to generate graph based on user query
 def generate_graph(data, query):
     """Generate a graph based on user query."""
@@ -40,6 +52,8 @@ def generate_graph(data, query):
         sales_col = infer_column(data, {"gross_sales", "sales", "revenue"})
         date_col = infer_column(data, {"date", "time"})
 
+        query_type = classify_query(query)
+
         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()