Update app.py

app.py

@@ -1,5 +1,4 @@
 # app.py
-import io
 import tempfile
 from datetime import datetime
 
@@ -12,7 +11,6 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 # Config
 # ----------------------------
 
-# 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",
@@ -27,13 +25,12 @@ _MODEL_CACHE = {}  # cache: model_id -> (tokenizer, model)
 # ----------------------------
 
 def _load_model(model_id: str):
-    """Load tokenizer and model (cached)."""
     if model_id in _MODEL_CACHE:
         return _MODEL_CACHE[model_id]
 
     tok = AutoTokenizer.from_pretrained(model_id, use_fast=True)
 
-    # Ensure we have a pad token to avoid generate() warnings/errors.
+    # Ensure pad token exists for generate()
     if tok.pad_token is None:
         if tok.eos_token is not None:
             tok.pad_token = tok.eos_token
@@ -47,8 +44,6 @@ def _load_model(model_id: str):
         low_cpu_mem_usage=True,
         device_map="auto",
     )
-
-    # If we added tokens, resize embeddings.
     if model.get_input_embeddings().num_embeddings != len(tok):
         model.resize_token_embeddings(len(tok))
 
@@ -57,7 +52,6 @@ def _load_model(model_id: str):
 
 
 def _format_prompt(tokenizer, system_prompt: str, user_prompt: str) -> str:
-    """Prefer each model's chat template; fallback to a simple instruction format."""
     sys = (system_prompt or "").strip()
     usr = (user_prompt or "").strip()
 
@@ -72,12 +66,11 @@ def _format_prompt(tokenizer, system_prompt: str, user_prompt: str) -> str:
             add_generation_prompt=True,
         )
 
-    # 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"
 
 
-def generate_batch(
+def generate_batch_df(
     model_id: str,
     system_prompt: str,
     prompts_multiline: str,
@@ -87,16 +80,13 @@ def generate_batch(
     top_k: int,
     repetition_penalty: float,
 ) -> pd.DataFrame:
-    """Generate responses for multiple user prompts (one per line)."""
     tok, model = _load_model(model_id)
     device = model.device
 
-    # 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 and encode
     formatted = [_format_prompt(tok, system_prompt, p) for p in prompts]
     enc = tok(
         formatted,
@@ -105,7 +95,6 @@ def generate_batch(
         truncation=True,
     ).to(device)
 
-    # True prompt lengths per row (ignore padding)
     prompt_lens = enc["attention_mask"].sum(dim=1)
 
     with torch.no_grad():
@@ -121,7 +110,6 @@ def generate_batch(
             pad_token_id=tok.pad_token_id,
         )
 
-    # Slice generated tokens using prompt lengths
     responses, tokens_out = [], []
     for i in range(gen.size(0)):
         start = int(prompt_lens[i].item())
@@ -135,9 +123,7 @@ def generate_batch(
     )
 
 
-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.
+def write_csv_path(df: pd.DataFrame) -> str:
     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)
@@ -153,7 +139,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** to save them.
+        Click **Generate** to get batched responses and a **downloadable CSV**.
         """
     )
 
@@ -186,9 +172,6 @@ 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 downloading
-            state_df = gr.State(value=None)
-
             out_df = gr.Dataframe(
                 headers=["user_prompt", "response", "tokens_out"],
                 datatype=["str", "str", "number"],
@@ -196,19 +179,16 @@ with gr.Blocks(title="Multi-Prompt Chat (System Prompt Control)") as demo:
                 wrap=True,
                 interactive=False,
                 row_count=(0, "dynamic"),
-                type="pandas",  # ensure callbacks receive a pandas DataFrame
+                type="pandas",  # ensures pandas goes into callbacks
             )
 
-            # 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")
+            # IMPORTANT: type="filepath" so we can return a string path
+            csv_out = gr.File(label="Scored CSV", interactive=False, type="filepath")
 
-    # -------- Callbacks --------
+    # -------- Callback: generate table AND CSV path in one go --------
 
     def _generate_cb(model_id, system_prompt, prompts_multiline, max_new_tokens, temperature, top_p, top_k, repetition_penalty):
-        df = generate_batch(
+        df = generate_batch_df(
             model_id=model_id,
             system_prompt=system_prompt,
             prompts_multiline=prompts_multiline,
@@ -218,23 +198,15 @@ 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  # (table, state)
+        csv_path = write_csv_path(df)
+        return df, csv_path  # DataFrame to table, path to File(type="filepath")
 
     run_btn.click(
         _generate_cb,
         inputs=[model_id, system_prompt, prompts_multiline, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
-        outputs=[out_df, state_df],
+        outputs=[out_df, csv_out],
         api_name="generate_batch",
     )
 
-    def _prepare_csv_cb(df_state):
-        # 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}])
-        path = write_csv_tempfile(df_state)
-        return gr.File.update(value=path, visible=True)
-
-    csv_btn.click(_prepare_csv_cb, inputs=[state_df], outputs=[out_file], api_name="download_csv")
-
 if __name__ == "__main__":
     demo.launch()