Update app.py

app.py

@@ -1,23 +1,26 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
 
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+# Use the specified model for inference
+client = InferenceClient("canstralian/rabbit-redux")
 
 def respond(message, history, system_message, max_tokens, temperature, top_p):
-    messages = [{"role": "system", "content": system_message}]
-    for val in history:
-        if val[0]: messages.append({"role": "user", "content": val[0]})
-        if val[1]: messages.append({"role": "assistant", "content": val[1]})
-    messages.append({"role": "user", "content": message})
+    # Construct a single prompt from the system message, history, and current user message
+    prompt = system_message + "\n"
+    for user_msg, assistant_msg in history:
+        prompt += f"User: {user_msg}\nAssistant: {assistant_msg}\n"
+    prompt += f"User: {message}\nAssistant:"
+
+    # Generate the response
     response = ""
-    for message in client.chat_completion(
-        messages, max_tokens=max_tokens, stream=True, temperature=temperature, top_p=top_p
+    for output in client.text_generation(
+        prompt, max_tokens=max_tokens, temperature=temperature, top_p=top_p, stream=True
     ):
-        token = message.choices[0].delta.content
+        token = output.token.text
         response += token
         yield response
 
-# Updated system message for penetration testing
+# Set up the Gradio interface
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[