Update src/streamlit_app.py

src/streamlit_app.py

@@ -6,6 +6,7 @@ import plotly.express as px
 import plotly.figure_factory as ff
 from dotenv import load_dotenv
 from huggingface_hub import InferenceClient, login
+from google import genai
 from io import StringIO
 
 # ======================================================
@@ -13,18 +14,26 @@ from io import StringIO
 # ======================================================
 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.")
+st.caption("AI that cleans, analyzes, and visualizes your data — Hugging Face + Gemini compatible.")
 
 # ======================================================
 # 🔐 Load Environment Variables
 # ======================================================
 load_dotenv()
 HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_API_KEY")
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+
 if not HF_TOKEN:
     st.error("❌ Missing HF_TOKEN. Please set it in your .env file.")
 else:
     login(token=HF_TOKEN)
 
+if GEMINI_API_KEY:
+    gemini_client = genai.Client(api_key=GEMINI_API_KEY)
+else:
+    gemini_client = None
+    st.warning("⚠️ Gemini API key missing. Gemini 2.5 Flash will not work.")
+
 # ======================================================
 # 🧠 MODEL SETUP
 # ======================================================
@@ -34,7 +43,7 @@ with st.sidebar:
     CLEANER_MODEL = st.selectbox(
         "Select Cleaner Model:",
         [
-            "Qwen/Qwen2.5-Coder-7B-Instruct",
+            "Qwen/Qwen2.5-Coder-14B",
             "meta-llama/Meta-Llama-3-8B-Instruct",
             "microsoft/Phi-3-mini-4k-instruct",
             "mistralai/Mistral-7B-Instruct-v0.3"
@@ -46,6 +55,7 @@ with st.sidebar:
         "Select Analysis Model:",
         [
             "Qwen/Qwen2.5-14B-Instruct",
+            "Gemini 2.5 Flash (Google)",
             "mistralai/Mistral-7B-Instruct-v0.3",
             "HuggingFaceH4/zephyr-7b-beta"
         ],
@@ -56,17 +66,16 @@ with st.sidebar:
     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)
+hf_cleaner_client = InferenceClient(model=CLEANER_MODEL, token=HF_TOKEN)
+hf_analyst_client = None
+if ANALYST_MODEL != "Gemini 2.5 Flash (Google)":
+    hf_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.
-    """
+    """HF text generation fallback to chat_completion"""
     try:
         resp = client.text_generation(
             prompt,
@@ -90,7 +99,6 @@ def safe_hf_generate(client, prompt, temperature=0.3, max_tokens=512):
 # 🧩 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]
@@ -107,8 +115,8 @@ def fallback_clean(df: pd.DataFrame) -> pd.DataFrame:
 
 
 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)
+    """Clean dataset using AI. Full dataset sent for thorough cleaning."""
+    csv_text = df.to_csv(index=False)
     prompt = f"""
 You are a professional data cleaning assistant.
 Clean and standardize the dataset below dynamically:
@@ -118,31 +126,19 @@ Clean and standardize the dataset below dynamically:
 4. Remove irrelevant or duplicate rows
 Return ONLY a valid CSV text (no markdown, no explanations).
 
---- RAW SAMPLE ---
-{raw_preview}
+Dataset:
+{csv_text}
 """
 
     try:
-        cleaned_str = safe_hf_generate(cleaner_client, prompt, temperature=0.1, max_tokens=1024)
+        cleaned_str = safe_hf_generate(hf_cleaner_client, prompt, temperature=0.1, max_tokens=4096)
     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)
-
+    cleaned_str = cleaned_str.replace("```csv", "").replace("```", "").replace("###", "").strip()
     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:
@@ -150,8 +146,10 @@ Return ONLY a valid CSV text (no markdown, no explanations).
         return fallback_clean(df)
 
 
+# ======================================================
+# 🧩 DATA ANALYSIS
+# ======================================================
 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())
@@ -167,18 +165,13 @@ def summarize_dataframe(df: pd.DataFrame) -> str:
 
 
 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)
+    csv_text = df.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}
+--- FULL DATA SAMPLE ---
+{csv_text}
 
 --- USER QUESTION ---
 {user_query}
@@ -189,13 +182,21 @@ Respond with:
 3. Notable relationships or anomalies
 4. Data-driven recommendations
 """
-
     try:
-        response = safe_hf_generate(analyst_client, prompt, temperature=temperature, max_tokens=max_tokens)
-        return response
+        if ANALYST_MODEL == "Gemini 2.5 Flash (Google)":
+            if gemini_client is None:
+                return "⚠️ Gemini API key missing."
+            response = gemini_client.models.generate_content(
+                model="gemini-2.5-flash",
+                contents=[prompt]
+            )
+            return getattr(response, "text", "No response from Gemini.")
+        else:
+            return safe_hf_generate(hf_analyst_client, prompt, temperature=temperature, max_tokens=max_tokens)
     except Exception as e:
         return f"⚠️ Analysis failed: {e}"
 
+
 # ======================================================
 # 🚀 MAIN APP LOGIC
 # ======================================================
@@ -264,5 +265,6 @@ if uploaded:
             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.")