Update app.py

app.py

@@ -1,5 +1,6 @@
 # app.py
 import io
+import tempfile
 from datetime import datetime
 
 import gradio as gr
@@ -11,26 +12,20 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 # Config
 # ----------------------------
 
-# Small, free, instruction-tuned models that run on CPU in a Basic Space.
+# Small, free, instruction-tuned models that can run on CPU (Basic Space).
 DEFAULT_MODELS = [
     "google/gemma-2-2b-it",
     "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
     "Qwen/Qwen2.5-1.5B-Instruct",
 ]
 
-_MODEL_CACHE = {}  # (tokenizer, model) cache
+_MODEL_CACHE = {}  # cache: model_id -> (tokenizer, model)
 
 
 # ----------------------------
 # Utilities
 # ----------------------------
 
-def df_to_csv_bytes(df: pd.DataFrame) -> bytes:
-    buf = io.StringIO()
-    df.to_csv(buf, index=False)
-    return buf.getvalue().encode("utf-8")
-
-
 def _load_model(model_id: str):
     """Load tokenizer and model (cached)."""
     if model_id in _MODEL_CACHE:
@@ -38,9 +33,8 @@ def _load_model(model_id: str):
 
     tok = AutoTokenizer.from_pretrained(model_id, use_fast=True)
 
-    # Ensure we have a pad token to avoid warnings in generate
+    # Ensure we have a pad token to avoid generate() warnings/errors.
     if tok.pad_token is None:
-        # Prefer eos_token, else add a pad token
         if tok.eos_token is not None:
             tok.pad_token = tok.eos_token
         else:
@@ -53,7 +47,8 @@ def _load_model(model_id: str):
         low_cpu_mem_usage=True,
         device_map="auto",
     )
-    # If we added a pad token, resize embeddings
+
+    # If we added tokens, resize embeddings.
     if model.get_input_embeddings().num_embeddings != len(tok):
         model.resize_token_embeddings(len(tok))
 
@@ -62,9 +57,7 @@ def _load_model(model_id: str):
 
 
 def _format_prompt(tokenizer, system_prompt: str, user_prompt: str) -> str:
-    """
-    Prefer the model's chat template. Fallback to a light instruction format.
-    """
+    """Prefer each model's chat template; fallback to a simple instruction format."""
     sys = (system_prompt or "").strip()
     usr = (user_prompt or "").strip()
 
@@ -79,7 +72,7 @@ def _format_prompt(tokenizer, system_prompt: str, user_prompt: str) -> str:
             add_generation_prompt=True,
         )
 
-    # Fallback format
+    # Fallback plain format
     prefix = f"<<SYS>>\n{sys}\n<</SYS>>\n\n" if sys else ""
     return f"{prefix}<<USER>>\n{usr}\n<</USER>>\n<<ASSISTANT>>\n"
 
@@ -98,14 +91,13 @@ def generate_batch(
     tok, model = _load_model(model_id)
     device = model.device
 
-    # Split lines, discard empties
+    # Split lines, drop empties
     prompts = [p.strip() for p in prompts_multiline.splitlines() if p.strip()]
     if not prompts:
         return pd.DataFrame([{"user_prompt": "", "response": "", "tokens_out": 0}])
 
-    # Build formatted prompts per model
+    # Build formatted prompts and encode
     formatted = [_format_prompt(tok, system_prompt, p) for p in prompts]
-
     enc = tok(
         formatted,
         return_tensors="pt",
@@ -113,7 +105,7 @@ def generate_batch(
         truncation=True,
     ).to(device)
 
-    # True prompt lengths per row (use attention mask sum to ignore padding)
+    # True prompt lengths per row (ignore padding)
     prompt_lens = enc["attention_mask"].sum(dim=1)
 
     with torch.no_grad():
@@ -129,9 +121,8 @@ def generate_batch(
             pad_token_id=tok.pad_token_id,
         )
 
-    # Slice generated tokens per row using actual prompt length
-    responses = []
-    tokens_out = []
+    # Slice generated tokens using prompt lengths
+    responses, tokens_out = [], []
     for i in range(gen.size(0)):
         start = int(prompt_lens[i].item())
         gen_ids = gen[i, start:]
@@ -139,14 +130,18 @@ def generate_batch(
         responses.append(text)
         tokens_out.append(len(gen_ids))
 
-    df = pd.DataFrame(
-        {
-            "user_prompt": prompts,
-            "response": responses,
-            "tokens_out": tokens_out,
-        }
+    return pd.DataFrame(
+        {"user_prompt": prompts, "response": responses, "tokens_out": tokens_out}
     )
-    return df
+
+
+def write_csv_tempfile(df: pd.DataFrame) -> str:
+    """Write CSV to a real temp file and return its path (works in Spaces)."""
+    # Use NamedTemporaryFile with delete=False so Gradio can read after returning.
+    ts = datetime.utcnow().strftime("%Y%m%d-%H%M%S")
+    tmp = tempfile.NamedTemporaryFile(prefix=f"batch_{ts}_", suffix=".csv", delete=False, dir="/tmp")
+    df.to_csv(tmp.name, index=False)
+    return tmp.name
 
 
 # ----------------------------
@@ -158,7 +153,7 @@ with gr.Blocks(title="Multi-Prompt Chat (System Prompt Control)") as demo:
         """
         # 🧪 Multi-Prompt Chat for HF Space
         Pick a small free model, set a **system prompt**, and enter **multiple user prompts** (one per line).
-        Click **Generate** to get batched responses, then **Download CSV** for offline use.
+        Click **Generate** to get batched responses, then **Download CSV** to save them.
         """
     )
 
@@ -191,7 +186,7 @@ with gr.Blocks(title="Multi-Prompt Chat (System Prompt Control)") as demo:
             run_btn = gr.Button("Generate", variant="primary")
 
         with gr.Column(scale=1):
-            # Keep last results for stable downloads
+            # Keep last results for downloading
             state_df = gr.State(value=None)
 
             out_df = gr.Dataframe(
@@ -201,14 +196,14 @@ with gr.Blocks(title="Multi-Prompt Chat (System Prompt Control)") as demo:
                 wrap=True,
                 interactive=False,
                 row_count=(0, "dynamic"),
-                type="pandas",  # ensure callbacks get a pandas DataFrame
+                type="pandas",  # ensure callbacks receive a pandas DataFrame
             )
 
-            # Older Gradio versions don't support file_name on DownloadButton
-            download_btn = gr.DownloadButton(
-                label="Download CSV",
-                value=None,  # we update this with bytes on demand
-            )
+            # File widget that will display a real downloadable file
+            out_file = gr.File(label="Download CSV", visible=False)
+
+            # Separate button to trigger file creation
+            csv_btn = gr.Button("Prepare CSV for download")
 
     # -------- Callbacks --------
 
@@ -223,7 +218,7 @@ with gr.Blocks(title="Multi-Prompt Chat (System Prompt Control)") as demo:
             top_k=int(top_k),
             repetition_penalty=float(repetition_penalty),
         )
-        return df, df  # show in table, also store in state
+        return df, df  # (table, state)
 
     run_btn.click(
         _generate_cb,
@@ -233,14 +228,13 @@ with gr.Blocks(title="Multi-Prompt Chat (System Prompt Control)") as demo:
     )
 
     def _prepare_csv_cb(df_state):
-        # Fallback-safe: produce bytes only (older Gradio uses a default filename)
+        # Robust across Gradio versions: write to a real temp file and return its path
         if df_state is None or len(df_state) == 0:
             df_state = pd.DataFrame([{"user_prompt": "", "response": "", "tokens_out": 0}])
-        csv_bytes = df_to_csv_bytes(df_state)
-        # Some Gradio versions ignore filename updates; return bytes only for compatibility
-        return gr.DownloadButton.update(value=csv_bytes)
+        path = write_csv_tempfile(df_state)
+        return gr.File.update(value=path, visible=True)
 
-    download_btn.click(_prepare_csv_cb, inputs=[state_df], outputs=[download_btn], api_name="download_csv")
+    csv_btn.click(_prepare_csv_cb, inputs=[state_df], outputs=[out_file], api_name="download_csv")
 
 if __name__ == "__main__":
     demo.launch()