Update app.py

app.py

@@ -1,146 +1,69 @@
-import pandas as pd
-import matplotlib.pyplot as plt
-import tempfile
-import gradio as gr
-from transformers import pipeline
-from sklearn.cluster import KMeans
-import numpy as np
-
-class TaskAgent:
-    def parse_query(self, query):
-        if "trend" in query.lower():
-            return {"task": "trend_analysis"}
-        elif "cluster" in query.lower():
-            return {"task": "clustering"}
-        else:
-            return {"task": "distribution"}
+!pip install gradio seaborn smolagents transformers -q -U
 
-class AnalysisAgent:
-    def generate_insights(self, data):
-        numeric_data = data.select_dtypes(include=np.number)
-        insights = [
-            f"Rows: {len(data)}",
-            f"Columns: {len(data.columns)}",
-            "Means: " + ", ".join([f"{col}: {val:.2f}" 
-                                for col, val in numeric_data.mean().items()])
-        ]
-        return "\n".join(insights)
+import gradio as gr
+from smolagents import HfApiModel, CodeAgent
+from huggingface_hub import login
+import os
+import shutil
 
-class DataAnalysisAgent:
-    def __init__(self):
-        self.nlp_pipeline = pipeline("text-classification", 
-                                   model="distilbert-base-uncased")
-        self.task_agent = TaskAgent()
-        self.analysis_agent = AnalysisAgent()
+# Authenticate with Hugging Face
+login(os.getenv("HF_TOKEN"))
 
-    def preprocess_data(self, data):
-        data = data.dropna(axis=1, how='all')
-        data = data.fillna(data.mean(numeric_only=True))
-        for col in data.columns:
-            if data[col].dtype == 'object':
-                try:
-                    data[col] = pd.to_datetime(data[col])
-                except:
-                    pass
-        return data
+def analyze_data(csv_file, additional_notes=""):
+    # Clear previous figures
+    if os.path.exists('./figures'):
+        shutil.rmtree('./figures')
+    os.makedirs('./figures', exist_ok=True)
 
-    def parse_query(self, query):
-        return self.task_agent.parse_query(query)
+    # Initialize model and agent
+    model = HfApiModel("meta-llama/Llama-3.1-70B-Instruct")
+    agent = CodeAgent(
+        tools=[],
+        model=model,
+        additional_authorized_imports=["numpy", "pandas", "matplotlib.pyplot", "seaborn"],
+        max_iterations=10,
+    )
 
-    def analyze(self, data, task):
-        if task['task'] == 'trend_analysis':
-            group_col = task['params'].get('group_by', data.columns[0])
-            return data.groupby(group_col).mean().reset_index()
-        elif task['task'] == 'clustering':
-            numeric_cols = data.select_dtypes(include=np.number).columns
-            kmeans = KMeans(n_clusters=3)
-            data['cluster'] = kmeans.fit_predict(data[numeric_cols])
-            return data
-        elif task['task'] == 'distribution':
-            return data.describe()
-        return data
+    # Run analysis
+    analysis_result = agent.run(
+        """You are an expert data analyst. Perform comprehensive analysis including:
+        1. Basic statistics and data quality checks
+        2. 3 insightful analytical questions about relationships in the data
+        3. Visualization of key patterns and correlations
+        4. Actionable real-world insights derived from findings
+        
+        Generate publication-quality visualizations and save to './figures/'
+        """,
+        additional_args={
+            "additional_notes": additional_notes,
+            "source_file": csv_file
+        }
+    )
 
-    def visualize(self, data, task):
-        try:
-            fig, ax = plt.subplots(figsize=(8, 4))
-            if task['task'] == 'trend_analysis':
-                data.plot(kind='line', ax=ax, title='Trend Analysis')
-            elif task['task'] == 'clustering':
-                ax.scatter(data.iloc[:, 0], data.iloc[:, 1], c=data['cluster'])
-                ax.set_title('Cluster Analysis')
-            else:
-                data.plot(kind='bar', ax=ax, title='Distribution Analysis')
-            
-            plt.tight_layout()
-            temp_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
-            plt.savefig(temp_file.name)
-            plt.close()
-            return temp_file.name
-        except Exception as e:
-            print(f"Visualization error: {e}")
-            return None
+    # Collect generated visuals
+    visuals = [os.path.join('./figures', f) for f in os.listdir('./figures') 
+              if f.endswith(('.png', '.jpg', '.jpeg'))]
 
-def process_data(file, query):
-    agent = DataAnalysisAgent()
-    
-    try:
-        # Validate CSV file
-        if not file.name.endswith('.csv'):
-            raise ValueError("Please upload a CSV file")
-            
-        df = pd.read_csv(file.name)
-        df = agent.preprocess_data(df)
-        
-        task = agent.parse_query(query)
-        analyzed_data = agent.analyze(df, task)
-        
-        insights = agent.analysis_agent.generate_insights(analyzed_data)
-        chart_path = agent.visualize(analyzed_data, task)
-        
-        return insights, chart_path if chart_path else None
-        
-    except Exception as e:
-        return f"Error: {str(e)}", None
+    return analysis_result, visuals
 
-# Gradio Interface with CSV-specific upload
+# Create Gradio interface
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("""
-    # Intelligent Data Analysis Agent
-    **Upload your CSV file and ask questions about your data**
-    """)
+    gr.Markdown("## 📊 AI Data Analysis Agent")
     
     with gr.Row():
         with gr.Column():
-            gr.Markdown("## Step 1: Upload your CSV file")
-            file_input = gr.File(
-                label="Select CSV File",
-                file_types=[".csv"],
-                type="filepath"
-            )
-            gr.Markdown("## Step 2: Ask your question")
-            query_input = gr.Textbox(
-                label="Data Question",
-                placeholder="e.g., 'Show sales trends by region'"
-            )
-            submit_btn = gr.Button("Analyze", variant="primary")
-            
+            file_input = gr.File(label="Upload CSV Dataset", type="filepath")
+            notes_input = gr.Textbox(label="Dataset Notes (Optional)", lines=3)
+            analyze_btn = gr.Button("Analyze", variant="primary")
+        
         with gr.Column():
-            gr.Markdown("## Analysis Results")
-            insights_output = gr.Textbox(
-                label="Insights",
-                interactive=False,
-                lines=10
-            )
-            chart_output = gr.Image(
-                label="Visualization",
-                show_label=True
-            )
-    
-    submit_btn.click(
-        fn=process_data,
-        inputs=[file_input, query_input],
-        outputs=[insights_output, chart_output]
+            analysis_output = gr.Textbox(label="Analysis Report", interactive=False)
+            gallery = gr.Gallery(label="Data Visualizations", columns=2)
+
+    analyze_btn.click(
+        fn=analyze_data,
+        inputs=[file_input, notes_input],
+        outputs=[analysis_output, gallery]
     )
 
-if __name__ == "__main__":
-    demo.queue().launch(debug=True)
+demo.launch(debug=True)