Update app.py

app.py

@@ -1,46 +1,60 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
 
-# Available open-source base models
+# Pick three open-source models (note: some may require login/accept T&Cs)
 MODEL_CHOICES = {
     "Mistral 7B Instruct (mistralai/Mistral-7B-Instruct-v0.2)": "mistralai/Mistral-7B-Instruct-v0.2",
     "Falcon 7B Instruct (tiiuae/falcon-7b-instruct)": "tiiuae/falcon-7b-instruct",
-    "LLaMA-2 7B Chat (meta-llama/Llama-2-7b-chat-hf)": "meta-llama/Llama-2-7b-chat-hf",
+    "Zephyr 7B Beta (HuggingFaceH4/zephyr-7b-beta)": "HuggingFaceH4/zephyr-7b-beta",
 }
 
-
-def complete_text(prompt, max_tokens, temperature, top_p, model_choice, hf_token: gr.OAuthToken):
+def complete_text(
+    prompt: str,
+    max_tokens: int,
+    temperature: float,
+    top_p: float,
+    model_choice: str,
+    hf_token: gr.OAuthToken,   # <-- token comes from LoginButton
+):
     """
-    Get a plain text completion from a Hugging Face-hosted open-source LLM.
-    Streams output token-by-token.
+    Stream a plain text completion from a HF-hosted open model.
     """
-    if not hf_token or not hf_token.token:
-        yield "⚠️ Please log in with your Hugging Face account (needed for gated models like LLaMA-2)."
+    if not prompt:
+        yield "⚠️ Enter a prompt."
         return
 
     model_id = MODEL_CHOICES[model_choice]
-    client = InferenceClient(model=model_id, token=hf_token.token)
 
-    response_text = ""
-    stream = client.text_generation(
-        prompt,
-        max_new_tokens=max_tokens,
-        temperature=temperature,
-        top_p=top_p,
-        stream=True,
-        repetition_penalty=1.0,
-    )
+    # Some models require auth; if the user isn't logged in, warn them.
+    token_str = getattr(hf_token, "token", None)
+    if token_str is None:
+        yield "🔐 This Space uses the Hugging Face Inference API. Please click **Login** (left sidebar) to authorize."
+        return
 
-    for chunk in stream:
-        response_text += chunk
-        yield response_text
+    client = InferenceClient(model=model_id, token=token_str)
+
+    try:
+        generated = ""
+        # text_generation returns an iterator of string chunks when stream=True
+        for chunk in client.text_generation(
+            prompt=prompt,
+            max_new_tokens=int(max_tokens),
+            temperature=float(temperature),
+            top_p=float(top_p),
+            stream=True,
+            repetition_penalty=1.0,
+        ):
+            generated += chunk
+            yield generated
+    except Exception as e:
+        yield f"❌ Inference error: {e}"
 
 
 with gr.Blocks() as demo:
-    gr.Markdown("## ✍️ Text Completion Demo with Open-Source Base LLMs")
+    gr.Markdown("## ✍️ Text Completion Demo (Hugging Face Inference API)")
     gr.Markdown(
-        "Pick a model hosted on Hugging Face, enter a prompt, adjust decoding parameters, "
-        "and watch the model complete your text."
+        "Pick a model, enter a prompt, and see a streamed completion. "
+        "Some models require you to **Login** (left sidebar) to access gated weights."
     )
 
     with gr.Row():
@@ -50,33 +64,27 @@ with gr.Blocks() as demo:
                 placeholder="Type the beginning of your text...",
                 lines=4,
             )
-            max_tokens = gr.Slider(
-                minimum=1, maximum=1024, value=100, step=1, label="Max tokens"
-            )
-            temperature = gr.Slider(
-                minimum=0.0, maximum=2.0, value=0.7, step=0.1, label="Temperature"
-            )
-            top_p = gr.Slider(
-                minimum=0.1, maximum=1.0, value=1.0, step=0.05, label="Top-p"
-            )
+            max_tokens = gr.Slider(1, 1024, value=100, step=1, label="Max new tokens")
+            temperature = gr.Slider(0.0, 2.0, value=0.7, step=0.1, label="Temperature")
+            top_p = gr.Slider(0.1, 1.0, value=1.0, step=0.05, label="Top-p")
             model_choice = gr.Dropdown(
                 choices=list(MODEL_CHOICES.keys()),
                 value=list(MODEL_CHOICES.keys())[0],
                 label="Choose a model",
             )
-            submit = gr.Button("Generate Completion")
+            submit = gr.Button("Generate")
         with gr.Column(scale=3):
-            output = gr.Textbox(
-                label="Generated Completion",
-                lines=15,
-            )
+            with gr.Sidebar():
+                login = gr.LoginButton()  # <-- THIS provides the OAuth token
+            output = gr.Textbox(label="Generated Completion", lines=15)
 
-    # ✅ Pass gr.OAuthToken as an input type (no manual instantiation)
+    # Wire the login button as the input for the token parameter
     submit.click(
         fn=complete_text,
-        inputs=[prompt, max_tokens, temperature, top_p, model_choice, gr.OAuthToken],
+        inputs=[prompt, max_tokens, temperature, top_p, model_choice, login],
         outputs=output,
     )
 
 if __name__ == "__main__":
     demo.launch()
+