Update app.py

app.py

@@ -4,43 +4,45 @@ import torch
 from transformers import pipeline
 from openai import OpenAI
 
-# ------------------------------------------------------------
-# Model registry
-# ------------------------------------------------------------
+# -------------------------
+# Model choices
+# -------------------------
 MODEL_OPTIONS = [
-    "GPT-1 (OpenAI GPT) - local",
-    "GPT-2 (small) - local",
-    "GPT-3 (text-davinci-003) - OpenAI",
+    "GPT-1 (openai-gpt) - local",
+    "GPT-2 (gpt2) - local",
     "GPT-3.5 (gpt-3.5-turbo) - OpenAI",
 ]
 
 MODEL_MAP = {
-    "GPT-1 (OpenAI GPT) - local": {"kind": "hf", "id": "openai-gpt"},
-    "GPT-2 (small) - local": {"kind": "hf", "id": "gpt2"},
-    "GPT-3 (text-davinci-003) - OpenAI": {"kind": "openai-completion", "id": "text-davinci-003"},
+    "GPT-1 (openai-gpt) - local": {"kind": "hf", "id": "openai-gpt"},
+    "GPT-2 (gpt2) - local": {"kind": "hf", "id": "gpt2"},
     "GPT-3.5 (gpt-3.5-turbo) - OpenAI": {"kind": "openai-chat", "id": "gpt-3.5-turbo"},
 }
 
-# Cache for local HF pipelines
+# Cache pipelines for HF models so we only load once
 HF_PIPELINES = {}
 
-# OpenAI client (only initialized if key exists)
+# OpenAI client (only if key exists)
 OPENAI_KEY = os.getenv("OPENAI_API_KEY")
 OPENAI_CLIENT = OpenAI(api_key=OPENAI_KEY) if OPENAI_KEY else None
 
 
 def get_hf_pipeline(model_id: str):
-    """Create or fetch a cached text-generation pipeline for a HF model."""
+    """Create/fetch a lightweight text-generation pipeline for CPU/GPU."""
     if model_id in HF_PIPELINES:
         return HF_PIPELINES[model_id]
 
     device = 0 if torch.cuda.is_available() else -1
-    gen = pipeline("text-generation", model=model_id, device=device)
+    gen = pipeline(
+        "text-generation",
+        model=model_id,
+        device=device,
+    )
     HF_PIPELINES[model_id] = gen
     return gen
 
 
-def generate(model_choice: str, prompt: str, max_new_tokens: int, temperature: float, top_p: float, seed: int):
+def generate(model_choice, prompt, max_new_tokens, temperature, top_p, seed):
     if not prompt.strip():
         return "Please enter a prompt."
 
@@ -48,46 +50,34 @@ def generate(model_choice: str, prompt: str, max_new_tokens: int, temperature: f
     kind = info["kind"]
     model_id = info["id"]
 
-    if seed is not None and seed >= 0:
-        torch.manual_seed(seed)
+    if seed is not None and int(seed) >= 0:
+        torch.manual_seed(int(seed))
 
     try:
         if kind == "hf":
             gen = get_hf_pipeline(model_id)
             out = gen(
                 prompt,
-                max_new_tokens=max_new_tokens,
+                max_new_tokens=int(max_new_tokens),
                 do_sample=temperature > 0,
-                temperature=max(1e-6, temperature),
-                top_p=top_p,
+                temperature=max(1e-6, float(temperature)),
+                top_p=float(top_p),
                 pad_token_id=gen.tokenizer.eos_token_id,
-                return_full_text=False,
+                return_full_text=False,  # don't echo the prompt
             )
             return out[0]["generated_text"]
 
-        if kind == "openai-completion":
-            if OPENAI_CLIENT is None:
-                return "⚠️ OPENAI_API_KEY not set. Add it in your Space secrets to use GPT-3."
-            resp = OPENAI_CLIENT.completions.create(
-                model=model_id,
-                prompt=prompt,
-                max_tokens=max_new_tokens,
-                temperature=temperature,
-                top_p=top_p,
-            )
-            return resp.choices[0].text.strip()
-
         if kind == "openai-chat":
             if OPENAI_CLIENT is None:
-                return "⚠️ OPENAI_API_KEY not set. Add it in your Space secrets to use GPT-3.5."
+                return "⚠️ To use GPT-3.5, set OPENAI_API_KEY in your Space (Settings → Variables & secrets)."
             resp = OPENAI_CLIENT.chat.completions.create(
                 model=model_id,
                 messages=[{"role": "user", "content": prompt}],
-                max_tokens=max_new_tokens,
-                temperature=temperature,
-                top_p=top_p,
+                max_tokens=int(max_new_tokens),
+                temperature=float(temperature),
+                top_p=float(top_p),
             )
-            return resp.choices[0].message.content.strip()
+            return (resp.choices[0].message.content or "").strip()
 
         return f"Unknown model kind: {kind}"
 
@@ -95,20 +85,22 @@ def generate(model_choice: str, prompt: str, max_new_tokens: int, temperature: f
         return f"❌ Error from {model_choice} ({model_id}): {str(e)}"
 
 
-def toggle_openai_visibility(choice):
-    """Show a helpful banner if OpenAI key is missing and user picked an OpenAI model."""
+def maybe_warn(choice):
+    """Show a small banner if user picked GPT-3.5 without an API key set."""
     info = MODEL_MAP[choice]
-    if "openai" in info["kind"] and OPENAI_CLIENT is None:
-        return gr.update(value="⚠️ To use GPT-3 / GPT-3.5, set OPENAI_API_KEY in your Space secrets.", visible=True)
+    needs_key = (info["kind"] == "openai-chat") and (OPENAI_CLIENT is None)
+    if needs_key:
+        return gr.update(value="⚠️ GPT-3.5 requires OPENAI_API_KEY in Space secrets.", visible=True)
     return gr.update(visible=False)
 
 
-with gr.Blocks(title="GPT Playground: GPT-1 / GPT-2 / GPT-3 / GPT-3.5") as demo:
+with gr.Blocks(title="Mini GPT Playground") as demo:
     gr.Markdown(
         """
-        # Simple GPT Playground
-        Type a prompt, pick a model, and generate a continuation or reply.  
-        **Local models** (GPT-1 / GPT-2) run with `transformers`. **OpenAI models** (GPT-3 / GPT-3.5) require `OPENAI_API_KEY`.
+        # Mini GPT Playground
+        Type a prompt and choose a model.  
+        **Local (HF):** GPT-1 / GPT-2 — runs in this Space container with `transformers`.  
+        **OpenAI (API):** GPT-3.5 — requires `OPENAI_API_KEY`.
         """
     )
 
@@ -118,29 +110,19 @@ with gr.Blocks(title="GPT Playground: GPT-1 / GPT-2 / GPT-3 / GPT-3.5") as demo:
     with gr.Row():
         temperature = gr.Slider(0.0, 2.0, value=0.8, step=0.05, label="Temperature")
         top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.01, label="Top-p")
-        seed = gr.Number(value=42, precision=0, label="Seed (set ≥0 to fix sampling)")
+        seed = gr.Number(value=42, precision=0, label="Seed (≥0 to fix sampling)")
 
     prompt = gr.Textbox(lines=6, label="Prompt", placeholder="Write a short story about a curious robot...")
-
-    warning_md = gr.Markdown("", visible=False)
+    warn = gr.Markdown("", visible=False)
 
     generate_btn = gr.Button("Generate", variant="primary")
     output = gr.Textbox(lines=12, label="Output")
 
-    model_choice.change(toggle_openai_visibility, inputs=[model_choice], outputs=[warning_md])
+    model_choice.change(maybe_warn, inputs=[model_choice], outputs=[warn])
     generate_btn.click(
         generate,
         inputs=[model_choice, prompt, max_new_tokens, temperature, top_p, seed],
         outputs=[output],
     )
 
-    gr.Markdown(
-        """
-        ---  
-        **Tips**
-        - If you see an error on GPT-3: it may no longer be enabled on your account. Try GPT-3.5.  
-        - Local models here are the small baseline versions (`openai-gpt`, `gpt2`) to keep Spaces lightweight.
-        """
-    )
-
 demo.queue(max_size=16).launch()