Update app.py

app.py

@@ -1,23 +1,24 @@
 from huggingface_hub import InferenceClient
 import gradio as gr
+import time
 
 client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
 
 # Your system prompt
-SYSTEM_PROMPT = "Your work is to rewrite and enhance the user input in sentence form without changing it's essence"
+SYSTEM_PROMPT = "Rewrite and enhance this in sentence form"
 
 def format_prompt(message, history=[]):
-    prompt = "<s>"  # Start sequence token
-    prompt += f"[SYSTEM] {SYSTEM_PROMPT} [/SYSTEM]"  # Add the system prompt here
-
+    prompt = "<s>"
+    prompt += f"[INST] SYSTEM: {SYSTEM_PROMPT} [/INST]"  # Add the system prompt here
     for user_prompt, bot_response in history:
-        prompt += f"[INST] {user_prompt} [/INST] {bot_response}"
-
-    prompt += f"[INST] {message} [/INST]</s>"  # Close the prompt sequence
+        prompt += f"[INST] {user_prompt} [/INST]"
+        prompt += f" {bot_response}</s> "
+    # Append the user message as the last entry in the prompt history
+    prompt += f"[INST] {message} [/INST]</s>"
     return prompt
 
 def generate(prompt, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0):
-    temperature = max(float(temperature), 1e-2)
+    temperature = float(temperature)
     top_p = float(top_p)
 
     generate_kwargs = dict(
@@ -29,13 +30,14 @@ def generate(prompt, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition
     )
 
     formatted_prompt = format_prompt(prompt)
+
     stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
     output = ""
 
     for response in stream:
         output += response.token.text
-        yield output.rstrip("</s>")
-    return output.rstrip("</s>")
+        yield output.strip('</s>')
+    return output.strip('</s>')
 
 with gr.Blocks() as demo:
     input_text = gr.Textbox(placeholder="Enter your prompt here...", lines=2, max_lines=2, label="Prompt")