Update app.py

app.py

@@ -1,17 +1,9 @@
 #!/usr/bin/env python3
 """
 ai_csv_editor_hf.py ── AI-powered CSV editor using a Hugging Face model on CPU.
-
-Features:
-- Upload one or more CSV files (main + optional lookup tables)
-- Type spreadsheet-style commands: CONCAT, VLOOKUP, XLOOKUP, SUMIF
-- LLM (google/flan-t5-base) converts commands → JSON “edit plan”
-- pandas applies each action in sequence
-- Preview first 20 rows & download modified CSV
 """
 
 import json
-import io
 import tempfile
 import textwrap
 import pathlib
@@ -31,14 +23,18 @@ TEMPERATURE  = 0.0
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model     = AutoModelForSeq2SeqLM.from_pretrained(
                 MODEL_NAME,
-                device_map="cpu",      # force CPU
+                device_map="cpu",      # force CPU placement
                 torch_dtype="auto"
             )
+
+# **Important change**: no `device=` argument here!
 generator = pipeline(
     "text2text-generation",
     model=model,
     tokenizer=tokenizer,
-    device=-1,                  # -1 = CPU
+    max_new_tokens=MAX_NEW_TOK,
+    temperature=TEMPERATURE,
+    do_sample=False,
 )
 
 # ──────────────────────────────────────────────────────────
@@ -76,12 +72,7 @@ into JSON edit plans. Respond with ONLY valid JSON matching this schema:
 
 def plan_from_command(cmd: str) -> Dict[str, Any]:
     prompt = f"{SYSTEM_PROMPT}\n\nUser: {cmd}\nJSON:"
-    output = generator(
-        prompt,
-        max_new_tokens=MAX_NEW_TOK,
-        temperature=TEMPERATURE,
-        do_sample=False,
-    )[0]["generated_text"]
+    output = generator(prompt)[0]["generated_text"]
     try:
         return json.loads(output)
     except json.JSONDecodeError as e:
@@ -105,7 +96,6 @@ def apply_action(df: pd.DataFrame,
 
     elif op in {"vlookup", "xlookup"}:
         lookup_df = uploads[act["lookup_file"]]
-        # select only the two relevant columns and rename for merging
         right = lookup_df[[act["lookup_column"], act["return_column"]]] \
             .rename(columns={
                 act["lookup_column"]: act["lookup_value"],
@@ -130,29 +120,24 @@ def run_editor(files: List[gr.File], command: str):
     if not files:
         return None, "⚠️ Please upload at least one CSV file.", None
 
-    # Load uploaded CSVs into a dictionary
     uploads = {
         pathlib.Path(f.name).name: pd.read_csv(f.name)
         for f in files
     }
-    # Treat the first file as the main dataset
     main_name = list(uploads.keys())[0]
     df = uploads[main_name]
 
-    # Generate plan
     try:
         plan = plan_from_command(command)
     except Exception as e:
         return None, f"❌ LLM error: {e}", None
 
-    # Apply actions
     try:
         for act in plan["actions"]:
             df = apply_action(df, uploads, act)
     except Exception as e:
         return None, f"❌ Execution error: {e}", None
 
-    # Write modified CSV to a temp file and return
     tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".csv")
     df.to_csv(tmp.name, index=False)
     return df.head(20), "✅ Success! Download below.", tmp.name
@@ -170,9 +155,9 @@ with gr.Blocks(title="AI CSV Editor (HF, CPU)") as demo:
     download  = gr.File(label="Download Result")
 
     run_btn.click(
-        fn=run_editor,
-        inputs=[csv_files, cmd_box],
-        outputs=[preview, status, download]
+        run_editor,
+        [csv_files, cmd_box],
+        [preview, status, download]
     )
 
 if __name__ == "__main__":