Update app.py

app.py

@@ -1,20 +1,35 @@
-import gradio as gr
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
-from threading import Thread
+# app.py
+import io
 from datetime import datetime
+
+import gradio as gr
 import pandas as pd
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# ---------- Config ----------
+# ----------------------------
+# Config
+# ----------------------------
 
-# Small, free chat models that run on CPU in a basic Space (pick one if you like)
+# Small, free, instruction-tuned models that run on CPU in a Basic Space.
 DEFAULT_MODELS = [
+    "google/gemma-2-2b-it",
     "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
     "Qwen/Qwen2.5-1.5B-Instruct",
 ]
 
-# Cache for loaded models to avoid reloading on each call
-_MODEL_CACHE = {}
+_MODEL_CACHE = {}  # (tokenizer, model) cache
+
+
+# ----------------------------
+# 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)."""
@@ -22,26 +37,36 @@ def _load_model(model_id: str):
         return _MODEL_CACHE[model_id]
 
     tok = AutoTokenizer.from_pretrained(model_id, use_fast=True)
-    # bfloat16 works on many CPUs and GPUs; fall back to float32 if needed
-    dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
 
+    # Ensure we have a pad token to avoid warnings in generate
+    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:
+            tok.add_special_tokens({"pad_token": "<|pad|>"})
+
+    dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
         torch_dtype=dtype,
         low_cpu_mem_usage=True,
         device_map="auto",
     )
+    # If we added a pad token, resize embeddings
+    if model.get_input_embeddings().num_embeddings != len(tok):
+        model.resize_token_embeddings(len(tok))
 
     _MODEL_CACHE[model_id] = (tok, model)
     return tok, model
 
+
 def _format_prompt(tokenizer, system_prompt: str, user_prompt: str) -> str:
     """
-    Use the model's chat template if available; otherwise
-    create a simple system+user concatenation.
+    Prefer the model's chat template. Fallback to a light instruction format.
     """
-    sys = system_prompt.strip() if system_prompt else ""
-    usr = user_prompt.strip()
+    sys = (system_prompt or "").strip()
+    usr = (user_prompt or "").strip()
 
     if hasattr(tokenizer, "apply_chat_template") and tokenizer.chat_template:
         messages = []
@@ -53,12 +78,11 @@ def _format_prompt(tokenizer, system_prompt: str, user_prompt: str) -> str:
             tokenize=False,
             add_generation_prompt=True,
         )
-    # Fallback: a lightweight instruction format
-    prompt = ""
-    if sys:
-        prompt += f"<<SYS>>\n{sys}\n<</SYS>>\n\n"
-    prompt += f"<<USER>>\n{usr}\n<</USER>>\n<<ASSISTANT>>\n"
-    return prompt
+
+    # Fallback 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(
     model_id: str,
@@ -69,72 +93,72 @@ def generate_batch(
     top_p: float,
     top_k: int,
     repetition_penalty: float,
-):
-    """Generate for multiple user prompts (one per line)."""
+) -> pd.DataFrame:
+    """Generate responses for multiple user prompts (one per line)."""
     tok, model = _load_model(model_id)
     device = model.device
 
-    # Split lines, drop empties
+    # Split lines, discard 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}])
 
-    # Prepare inputs
+    # Build formatted prompts per model
     formatted = [_format_prompt(tok, system_prompt, p) for p in prompts]
-    inputs = tok(
+
+    enc = tok(
         formatted,
         return_tensors="pt",
         padding=True,
         truncation=True,
     ).to(device)
 
+    # True prompt lengths per row (use attention mask sum to ignore padding)
+    prompt_lens = enc["attention_mask"].sum(dim=1)
+
     with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=max_new_tokens,
+        gen = model.generate(
+            **enc,
+            max_new_tokens=int(max_new_tokens),
             do_sample=(temperature > 0.0),
-            temperature=temperature if temperature > 0 else None,
-            top_p=top_p,
-            top_k=top_k if top_k > 0 else None,
-            repetition_penalty=repetition_penalty,
+            temperature=float(temperature) if temperature > 0 else None,
+            top_p=float(top_p),
+            top_k=int(top_k) if int(top_k) > 0 else None,
+            repetition_penalty=float(repetition_penalty),
             eos_token_id=tok.eos_token_id,
-            pad_token_id=tok.eos_token_id,
+            pad_token_id=tok.pad_token_id,
         )
 
-    # Slice off the prompt tokens to get only the generated text
-    gen_texts = []
-    for i in range(outputs.size(0)):
-        prompt_len = inputs["input_ids"][i].size(0)
-        # Some tokenizers need special handling; safest: decode full and strip prompt
-        full = tok.decode(outputs[i], skip_special_tokens=True)
-        prompt_only = tok.decode(inputs["input_ids"][i], skip_special_tokens=True)
-        # Remove the first occurrence of the prompt text
-        resp = full[len(prompt_only):].strip()
-        gen_texts.append(resp)
+    # Slice generated tokens per row using actual prompt length
+    responses = []
+    tokens_out = []
+    for i in range(gen.size(0)):
+        start = int(prompt_lens[i].item())
+        gen_ids = gen[i, start:]
+        text = tok.decode(gen_ids, skip_special_tokens=True).strip()
+        responses.append(text)
+        tokens_out.append(len(gen_ids))
 
     df = pd.DataFrame(
         {
             "user_prompt": prompts,
-            "response": gen_texts,
-            "tokens_out": [len(tok.encode(t)) for t in gen_texts],
+            "response": responses,
+            "tokens_out": tokens_out,
         }
     )
     return df
 
-def to_csv(df: pd.DataFrame):
-    ts = datetime.utcnow().strftime("%Y%m%d-%H%M%S")
-    path = f"/tmp/batch_{ts}.csv"
-    df.to_csv(path, index=False)
-    return path
 
-# ---------- UI ----------
+# ----------------------------
+# Gradio UI
+# ----------------------------
 
 with gr.Blocks(title="Multi-Prompt Chat (System Prompt Control)") as demo:
     gr.Markdown(
         """
-        # 🧪 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 as a table (downloadable as CSV).
+        # 🧪 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.
         """
     )
 
@@ -153,7 +177,7 @@ with gr.Blocks(title="Multi-Prompt Chat (System Prompt Control)") as demo:
             )
             prompts_multiline = gr.Textbox(
                 label="User prompts (one per line)",
-                placeholder="Write one query per line.\nExample:\nExplain transformers in simple terms\nGive 3 eco-friendly tips for students\nSummarise the benefits of multilingual models",
+                placeholder="One query per line.\nExample:\nExplain transformers in simple terms\nGive 3 eco-friendly tips for students\nSummarise the benefits of multilingual models",
                 lines=10,
             )
 
@@ -161,13 +185,15 @@ with gr.Blocks(title="Multi-Prompt Chat (System Prompt Control)") as demo:
                 max_new_tokens = gr.Slider(16, 1024, value=256, step=1, label="max_new_tokens")
                 temperature = gr.Slider(0.0, 2.0, value=0.7, step=0.05, label="temperature")
                 top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="top_p")
-                top_k = gr.Slider(0, 200, value=40, step=1, label="top_k (0 to disable)")
+                top_k = gr.Slider(0, 200, value=40, step=1, label="top_k (0 disables)")
                 repetition_penalty = gr.Slider(1.0, 2.0, value=1.1, step=0.01, label="repetition_penalty")
 
             run_btn = gr.Button("Generate", variant="primary")
-            csv_btn = gr.Button("Download CSV")
 
         with gr.Column(scale=1):
+            # Keep last results for stable downloads
+            state_df = gr.State(value=None)
+
             out_df = gr.Dataframe(
                 headers=["user_prompt", "response", "tokens_out"],
                 datatype=["str", "str", "number"],
@@ -175,11 +201,18 @@ with gr.Blocks(title="Multi-Prompt Chat (System Prompt Control)") as demo:
                 wrap=True,
                 interactive=False,
                 row_count=(0, "dynamic"),
-                type="pandas",
+                type="pandas",  # ensure callbacks get a pandas DataFrame
+            )
+
+            download_btn = gr.DownloadButton(
+                label="Download CSV",
+                value=None,                 # we update this with bytes on demand
+                file_name="batch.csv",
             )
-            out_file = gr.File(label="CSV file", visible=False)
 
-    def _generate(model_id, system_prompt, prompts_multiline, max_new_tokens, temperature, top_p, top_k, repetition_penalty):
+    # -------- Callbacks --------
+
+    def _generate_cb(model_id, system_prompt, prompts_multiline, max_new_tokens, temperature, top_p, top_k, repetition_penalty):
         df = generate_batch(
             model_id=model_id,
             system_prompt=system_prompt,
@@ -190,20 +223,23 @@ with gr.Blocks(title="Multi-Prompt Chat (System Prompt Control)") as demo:
             top_k=int(top_k),
             repetition_penalty=float(repetition_penalty),
         )
-        return df
+        return df, df  # show in table, also store in state
 
-    def _download(df):
-        path = to_csv(df)
-        return gr.File.update(value=path, visible=True)
-    
     run_btn.click(
-        _generate,
+        _generate_cb,
         inputs=[model_id, system_prompt, prompts_multiline, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
-        outputs=out_df,
+        outputs=[out_df, state_df],
         api_name="generate_batch",
     )
 
-    csv_btn.click(_download, inputs=out_df, outputs=out_file, api_name="download_csv")
+    def _prepare_csv_cb(df_state):
+        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)
+        ts = datetime.utcnow().strftime("%Y%m%d-%H%M%S")
+        return gr.DownloadButton.update(value=csv_bytes, file_name=f"batch_{ts}.csv")
+
+    download_btn.click(_prepare_csv_cb, inputs=[state_df], outputs=[download_btn], api_name="download_csv")
 
 if __name__ == "__main__":
     demo.launch()