Update app.py

app.py

@@ -1,38 +1,23 @@
-import streamlit as st
-import pandas as pd
 import os
-from csv_agent import generate_code
+from huggingface_hub import InferenceClient
 
-st.set_page_config(page_title="CSV Assistant with Hugging Face", layout="wide")
-st.title("🤖 CSV Assistant: Powered by Starcoder (Hugging Face)")
+# Initialize Hugging Face client
+client = InferenceClient(
+    model="Salesforce/codegen-2B-multi",
+    token=os.getenv("HUGGINGFACEHUB_API_TOKEN")
+)
 
-# Check for API token
-if os.getenv("HUGGINGFACEHUB_API_TOKEN") is None:
-    st.error("❌ API Token missing! Please set HUGGINGFACEHUB_API_TOKEN in your Space Secrets.")
-    st.stop()
+def generate_code(prompt):
+    full_prompt = f"""You are a professional Python Data Scientist.
+Given the following instruction, write Python pandas code to perform the task on a dataframe named df.
+Instruction: {prompt}
+Output only valid Python code. No explanations."""
 
-st.sidebar.title("Upload Your CSV")
-uploaded_file = st.sidebar.file_uploader("Choose a CSV file", type=["csv"])
-
-if uploaded_file:
-    df = pd.read_csv(uploaded_file)
-    st.write("### Preview of your data:", df.head())
-
-    prompt = st.text_input("💬 What would you like to do with the data?")
-    if prompt:
-        with st.spinner("Generating Python code using AI..."):
-            code = generate_code(prompt)
-            st.code(code, language="python")
-            
-            try:
-                exec_globals = {"df": df}
-                exec(code, exec_globals)
-                df = exec_globals["df"]
-                st.success("✅ Data updated successfully!")
-                st.write(df.head())
-
-                # Allow download
-                csv = df.to_csv(index=False).encode('utf-8')
-                st.download_button("⬇️ Download Updated CSV", csv, "updated_data.csv", "text/csv")
-            except Exception as e:
-                st.error(f"⚠️ Error while executing AI-generated code: {e}")
+    response = client.text_generation(
+        prompt=full_prompt,
+        max_new_tokens=300,
+        temperature=0.2,
+        top_p=0.95,
+        repetition_penalty=1.1
+    )
+    return response