kali-pentest / app.py
NexusInstruments's picture
Update app.py
2daa659 verified
Raw
History Blame Contribute Delete
3.3 kB
import gradio as gr
from huggingface_hub import InferenceClient
DEFAULT_MODEL = "openai/gpt-oss-20b"
DEFAULT_SYSTEM_MESSAGE = "You are a friendly Chatbot."
def respond(
message: str,
history: list[dict[str, str]],
system_message: str,
max_tokens: int,
temperature: float,
top_p: float,
hf_token: gr.OAuthToken,
):
"""
Chat completion handler with streaming, safe auth checks,
and graceful error handling.
"""
# --- Auth guard ---------------------------------------------------------
if hf_token is None or not hf_token.token:
yield "πŸ”’ **Authentication required.** Please log in using the sidebar button."
return
if not message or not message.strip():
yield "⚠️ Please enter a message before sending."
return
# --- Build messages -----------------------------------------------------
messages = []
if system_message and system_message.strip():
messages.append({"role": "system", "content": system_message})
for entry in history or []:
# Defensive normalisation: handle both dict and legacy tuple formats.
if isinstance(entry, dict) and "role" in entry and "content" in entry:
messages.append(entry)
elif isinstance(entry, (list, tuple)) and len(entry) >= 2:
user_msg, assistant_msg = str(entry[0]), str(entry[1])
messages.append({"role": "user", "content": user_msg})
if assistant_msg:
messages.append({"role": "assistant", "content": assistant_msg})
messages.append({"role": "user", "content": message})
# --- Stream inference ---------------------------------------------------
try:
client = InferenceClient(token=hf_token.token, model=DEFAULT_MODEL)
stream = client.chat_completion(
messages,
max_tokens=max_tokens,
stream=True,
temperature=temperature,
top_p=top_p,
)
response = ""
for chunk in stream:
choices = chunk.choices
if choices and choices[0].delta and choices[0].delta.content:
response += choices[0].delta.content
yield response
except Exception as e:
yield f"❌ **Inference error:** `{type(e).__name__}: {e}`"
# --- UI -------------------------------------------------------------------
chatbot = gr.ChatInterface(
respond,
type="messages", # Enforce the new {role, content} format
additional_inputs=[
gr.Textbox(
value=DEFAULT_SYSTEM_MESSAGE,
label="System message",
placeholder="You are a helpful assistant...",
),
gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
gr.Slider(
minimum=0.1,
maximum=1.0,
value=0.95,
step=0.05,
label="Top-p (nucleus sampling)",
),
],
)
with gr.Blocks() as demo:
with gr.Sidebar():
gr.Markdown("## πŸ” Authentication")
gr.LoginButton()
gr.LogoutButton()
chatbot.render()
if __name__ == "__main__":
demo.queue(default_concurrency_limit=20).launch()