Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,268 +1,112 @@
 import os
+import time
 import pandas as pd
-import numpy as np
 import streamlit as st
-import plotly.express as px
-import plotly.figure_factory as ff
+from io import StringIO
 from dotenv import load_dotenv
 from huggingface_hub import InferenceClient, login
-from io import StringIO
-
-# ======================================================
-# ⚙️ APP CONFIGURATION
-# ======================================================
-st.set_page_config(page_title="📊 Smart Data Analyst Pro", layout="wide")
-st.title("📊 Smart Data Analyst Pro")
-st.caption("AI that cleans, analyzes, and visualizes your data — powered by Hugging Face Inference API.")
 
-# ======================================================
-# 🔐 Load Environment Variables
-# ======================================================
+# ==========================================================
+# 🔐 Load environment + authenticate
+# ==========================================================
 load_dotenv()
-HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_API_KEY")
+HF_TOKEN = os.getenv("HF_TOKEN")
+
 if not HF_TOKEN:
-    st.error("❌ Missing HF_TOKEN. Please set it in your .env file.")
+    st.error("❌ Missing Hugging Face token. Please set HF_TOKEN in your .env file.")
 else:
     login(token=HF_TOKEN)
 
-# ======================================================
-# 🧠 MODEL SETUP
-# ======================================================
-with st.sidebar:
-    st.header("⚙️ Model Settings")
-
-    CLEANER_MODEL = st.selectbox(
-        "Select Cleaner Model:",
-        [
-            "Qwen/Qwen2.5-Coder-7B-Instruct",
-            "meta-llama/Meta-Llama-3-8B-Instruct",
-            "microsoft/Phi-3-mini-4k-instruct",
-            "mistralai/Mistral-7B-Instruct-v0.3"
-        ],
-        index=0
-    )
-
-    ANALYST_MODEL = st.selectbox(
-        "Select Analysis Model:",
-        [
-            "Qwen/Qwen2.5-14B-Instruct",
-            "mistralai/Mistral-7B-Instruct-v0.3",
-            "HuggingFaceH4/zephyr-7b-beta"
-        ],
-        index=0
-    )
-
-    temperature = st.slider("Temperature", 0.0, 1.0, 0.3)
-    max_tokens = st.slider("Max Tokens", 128, 2048, 512)
-
-# Initialize inference clients
-cleaner_client = InferenceClient(model=CLEANER_MODEL, token=HF_TOKEN)
-analyst_client = InferenceClient(model=ANALYST_MODEL, token=HF_TOKEN)
-
-# ======================================================
-# 🧩 SAFE GENERATION FUNCTION
-# ======================================================
-def safe_hf_generate(client, prompt, temperature=0.3, max_tokens=512):
-    """
-    Tries text_generation first, then falls back to chat_completion if not supported.
-    Returns plain string content.
-    """
-    try:
-        resp = client.text_generation(
-            prompt,
-            temperature=temperature,
-            max_new_tokens=max_tokens,
-            return_full_text=False,
-        )
-        return resp.strip()
-    except Exception as e:
-        if "Supported task: conversational" in str(e) or "not supported" in str(e):
-            chat_resp = client.chat_completion(
-                messages=[{"role": "user", "content": prompt}],
-                max_tokens=max_tokens,
-                temperature=temperature,
-            )
-            return chat_resp["choices"][0]["message"]["content"].strip()
-        else:
-            raise e
-
-# ======================================================
-# 🧩 SMART DATA CLEANING
-# ======================================================
-def fallback_clean(df: pd.DataFrame) -> pd.DataFrame:
-    """Backup rule-based cleaner."""
-    df = df.copy()
-    df.dropna(axis=1, how="all", inplace=True)
-    df.columns = [c.strip().replace(" ", "_").lower() for c in df.columns]
-    for col in df.columns:
-        if df[col].dtype == "O":
-            if not df[col].mode().empty:
-                df[col].fillna(df[col].mode()[0], inplace=True)
-            else:
-                df[col].fillna("Unknown", inplace=True)
-        else:
-            df[col].fillna(df[col].median(), inplace=True)
-    df.drop_duplicates(inplace=True)
-    return df
-
-
-def ai_clean_dataset(df: pd.DataFrame) -> pd.DataFrame:
-    """Cleans the dataset using the selected AI model. Falls back gracefully if the model fails."""
-    raw_preview = df.head(5).to_csv(index=False)
-    prompt = f"""
-You are a professional data cleaning assistant.
-Clean and standardize the dataset below dynamically:
-1. Handle missing values
-2. Fix column name inconsistencies
-3. Convert data types (dates, numbers, categories)
-4. Remove irrelevant or duplicate rows
-Return ONLY a valid CSV text (no markdown, no explanations).
-
---- RAW SAMPLE ---
-{raw_preview}
-"""
-
-    try:
-        cleaned_str = safe_hf_generate(cleaner_client, prompt, temperature=0.1, max_tokens=1024)
-    except Exception as e:
-        st.warning(f"⚠️ AI cleaning failed: {e}")
-        return fallback_clean(df)
-
-    cleaned_str = (
-        cleaned_str.replace("```csv", "")
-        .replace("```", "")
-        .replace("###", "")
-        .replace(";", ",")
-        .strip()
-    )
-
-    lines = cleaned_str.splitlines()
-    lines = [line for line in lines if "," in line and not line.lower().startswith(("note", "summary"))]
-    cleaned_str = "\n".join(lines)
-
-    try:
-        cleaned_df = pd.read_csv(StringIO(cleaned_str), on_bad_lines="skip")
-        cleaned_df = cleaned_df.dropna(axis=1, how="all")
-        cleaned_df.columns = [c.strip().replace(" ", "_").lower() for c in cleaned_df.columns]
-        return cleaned_df
-    except Exception as e:
-        st.warning(f"⚠️ AI CSV parse failed: {e}")
-        return fallback_clean(df)
-
-
-def summarize_dataframe(df: pd.DataFrame) -> str:
-    """Generate a concise summary of the dataframe."""
-    lines = [f"Rows: {len(df)} | Columns: {len(df.columns)}", "Column summaries:"]
-    for col in df.columns[:10]:
-        non_null = int(df[col].notnull().sum())
-        if pd.api.types.is_numeric_dtype(df[col]):
-            desc = df[col].describe().to_dict()
-            mean = float(desc.get("mean", np.nan))
-            median = float(df[col].median()) if non_null > 0 else None
-            lines.append(f"- {col}: mean={mean:.3f}, median={median}, non_null={non_null}")
-        else:
-            top = df[col].value_counts().head(3).to_dict()
-            lines.append(f"- {col}: top_values={top}, non_null={non_null}")
-    return "\n".join(lines)
-
-
-def query_analysis_model(df: pd.DataFrame, user_query: str, dataset_name: str) -> str:
-    """Send the dataframe and user query to the analysis model for interpretation."""
-    df_summary = summarize_dataframe(df)
-    sample = df.head(6).to_csv(index=False)
-    prompt = f"""
-You are a professional data analyst.
-Analyze the dataset '{dataset_name}' and answer the user's question.
-
---- SUMMARY ---
-{df_summary}
-
---- SAMPLE DATA ---
-{sample}
-
---- USER QUESTION ---
+# Create HF clients
+cleaner_client = InferenceClient(model="Qwen/Qwen2.5-Coder-14B", token=HF_TOKEN)
+analyst_client = InferenceClient(model="Qwen/Qwen2.5-14B-Instruct", token=HF_TOKEN)
+
+# ==========================================================
+# 🎛️ App Layout
+# ==========================================================
+st.set_page_config(page_title="🧹 Smart Data Analysis", page_icon="📊", layout="wide")
+st.title("📊 Smart Data Analysis Assistant")
+st.caption("Clean messy data, then run AI-powered insights and statistical analysis — all locally with open-source models.")
+
+# ==========================================================
+# 📁 Upload CSV
+# ==========================================================
+uploaded_file = st.file_uploader("📤 Upload your CSV dataset", type=["csv"])
+if uploaded_file:
+    df_raw = pd.read_csv(uploaded_file)
+    st.subheader("📄 Raw Data Preview")
+    st.dataframe(df_raw.head())
+
+    # ==========================================================
+    # 🧹 Data Cleaning
+    # ==========================================================
+    if st.button("🧹 Clean Data using Qwen Coder 14B"):
+        with st.spinner("Cleaning data... please wait ⏳"):
+            try:
+                # Convert DataFrame to text for cleaning
+                csv_text = df_raw.to_csv(index=False)
+
+                prompt = f"""
+You are a Python data cleaning assistant. 
+Clean this dataset and fix inconsistent column names, missing values, and formatting.
+Return a clean CSV version that can be loaded into pandas directly.
+
+Dataset:
+{csv_text}
+                """
+
+                response = cleaner_client.text_generation(
+                    prompt,
+                    temperature=0.2,
+                    max_new_tokens=2048,
+                )
+
+                cleaned_csv = response.strip().split("```")[-1]  # extract text
+                df_cleaned = pd.read_csv(StringIO(cleaned_csv))
+
+                st.session_state.cleaned_df = df_cleaned
+                st.success("✅ Data cleaned successfully!")
+                st.dataframe(df_cleaned.head())
+
+            except Exception as e:
+                st.error(f"⚠️ Cleaning failed: {e}")
+
+# ==========================================================
+# 📊 Data Analysis
+# ==========================================================
+if "cleaned_df" in st.session_state:
+    df = st.session_state.cleaned_df
+    st.divider()
+    st.subheader("📈 AI Data Analysis")
+
+    user_query = st.text_area("Ask about your data:", placeholder="e.g., What is the correlation between experience and salary?")
+    if st.button("🔍 Analyze"):
+        with st.spinner("Analyzing with Qwen 14B Instruct..."):
+            try:
+                csv_excerpt = df.head(30).to_csv(index=False)
+                analysis_prompt = f"""
+You are a data analyst. Analyze this dataset and answer the question.
+
+Data sample (CSV):
+{csv_excerpt}
+
+Question:
 {user_query}
 
-Respond with:
-1. Key insights and patterns
-2. Quantitative findings
-3. Notable relationships or anomalies
-4. Data-driven recommendations
+Instructions:
+- Be accurate and concise.
+- If numerical analysis is relevant, describe it.
+- Use markdown for readability.
 """
 
-    try:
-        response = safe_hf_generate(analyst_client, prompt, temperature=temperature, max_tokens=max_tokens)
-        return response
-    except Exception as e:
-        return f"⚠️ Analysis failed: {e}"
-
-# ======================================================
-# 🚀 MAIN APP LOGIC
-# ======================================================
-uploaded = st.file_uploader("📎 Upload CSV or Excel file", type=["csv", "xlsx"])
-
-if uploaded:
-    df = pd.read_csv(uploaded) if uploaded.name.endswith(".csv") else pd.read_excel(uploaded)
-
-    with st.spinner("🧼 AI Cleaning your dataset..."):
-        cleaned_df = ai_clean_dataset(df)
+                response = analyst_client.text_generation(
+                    analysis_prompt,
+                    temperature=0.5,
+                    max_new_tokens=1024,
+                )
 
-    st.subheader("✅ Cleaned Dataset Preview")
-    st.dataframe(cleaned_df.head(), use_container_width=True)
+                st.markdown("### 🧠 Analysis Result")
+                st.write(response.strip())
 
-    with st.expander("📋 Cleaning Summary", expanded=False):
-        st.text(summarize_dataframe(cleaned_df))
-
-    with st.expander("📈 Quick Visualizations", expanded=True):
-        numeric_cols = cleaned_df.select_dtypes(include="number").columns.tolist()
-        categorical_cols = cleaned_df.select_dtypes(exclude="number").columns.tolist()
-
-        viz_type = st.selectbox(
-            "Visualization Type",
-            ["Scatter Plot", "Histogram", "Box Plot", "Correlation Heatmap", "Categorical Count"]
-        )
-
-        if viz_type == "Scatter Plot" and len(numeric_cols) >= 2:
-            x = st.selectbox("X-axis", numeric_cols)
-            y = st.selectbox("Y-axis", numeric_cols, index=min(1, len(numeric_cols)-1))
-            color = st.selectbox("Color", ["None"] + categorical_cols)
-            fig = px.scatter(cleaned_df, x=x, y=y, color=None if color=="None" else color)
-            st.plotly_chart(fig, use_container_width=True)
-
-        elif viz_type == "Histogram" and numeric_cols:
-            col = st.selectbox("Column", numeric_cols)
-            fig = px.histogram(cleaned_df, x=col, nbins=30)
-            st.plotly_chart(fig, use_container_width=True)
-
-        elif viz_type == "Box Plot" and numeric_cols:
-            col = st.selectbox("Column", numeric_cols)
-            fig = px.box(cleaned_df, y=col)
-            st.plotly_chart(fig, use_container_width=True)
-
-        elif viz_type == "Correlation Heatmap" and len(numeric_cols) > 1:
-            corr = cleaned_df[numeric_cols].corr()
-            fig = ff.create_annotated_heatmap(
-                z=corr.values,
-                x=list(corr.columns),
-                y=list(corr.index),
-                annotation_text=corr.round(2).values,
-                showscale=True
-            )
-            st.plotly_chart(fig, use_container_width=True)
-
-        elif viz_type == "Categorical Count" and categorical_cols:
-            cat = st.selectbox("Category", categorical_cols)
-            fig = px.bar(cleaned_df[cat].value_counts().reset_index(), x="index", y=cat)
-            st.plotly_chart(fig, use_container_width=True)
-        else:
-            st.warning("⚠️ Not enough columns for this visualization type.")
-
-    st.subheader("💬 Ask AI About Your Data")
-    user_query = st.text_area("Enter your question:", placeholder="e.g. What factors influence sales the most?")
-    if st.button("Analyze with AI", use_container_width=True) and user_query:
-        with st.spinner("🤖 Interpreting data..."):
-            result = query_analysis_model(cleaned_df, user_query, uploaded.name)
-        st.markdown("### 💡 Insights")
-        st.markdown(result)
-else:
-    st.info("📥 Upload a dataset to begin smart analysis.")
+            except Exception as e:
+                st.error(f"⚠️ Analysis failed: {e}")