Update app.py

app.py

@@ -1,25 +1,26 @@
-import os
 import random
 import gradio as gr
 from huggingface_hub import InferenceClient
 
-hf_token = os.getenv("HF_TOKEN")
-client = InferenceClient("google/gemma-7b", use_auth_token=hf_token)
+
+client = InferenceClient("google/gemma-7b")
 
 def format_prompt(message, history):
     prompt = ""
     if history:
         for user_prompt, bot_response in history:
             prompt += f"<start_of_turn>user{user_prompt}<end_of_turn>"
-            prompt += f"<start_of_turn>model{bot_response}<end_of_turn>"
+            prompt += f"<start_of_turn>model{bot_response}"
     prompt += f"<start_of_turn>user{message}<end_of_turn><start_of_turn>model"
     return prompt
 
-def generate(prompt, history, temperature=0.7, max_new_tokens=1024, top_p=0.90, repetition_penalty=1.0):
-    # Adjust parameters as necessary
+
+def generate(prompt, history, temperature=0.7, max_new_tokens=1024, top_p=0.90, repetition_penalty=0.9):
     temperature = float(temperature)
+    if temperature < 1e-2:
+        temperature = 1e-2
     top_p = float(top_p)
-
+   
     if not history:
         history = []
 
@@ -36,34 +37,67 @@ def generate(prompt, history, temperature=0.7, max_new_tokens=1024, top_p=0.90,
 
     formatted_prompt = format_prompt(prompt, history)
 
-    stream = client(text=formatted_prompt, parameters=generate_kwargs, wait_for_model=True)
+    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
     output = ""
 
-    for response in stream["generated_text"]:
-        output += response
+    for response in stream:
+        output += response.token.text
         yield output
     history.append((prompt, output))
     return output
 
-# Setup Gradio Interface
-chatbot_ui = gr.Chatbot()
-
-def chat_interface(prompt, temperature=0.7, max_new_tokens=160, top_p=0.90, repetition_penalty=1.0):
-    history = []  # Initialize or fetch existing history
-    return generate(prompt, history, temperature, max_new_tokens, top_p, repetition_penalty)
+    
+mychatbot = gr.Chatbot(
+    avatar_images=["./user.png", "./botgm.png"], bubble_full_width=False, show_label=False, show_copy_button=True, likeable=True,)
 
-iface = gr.Interface(fn=chat_interface, 
-                     inputs=[gr.Textbox(label="Your Message"), 
-                             gr.Slider(label="Temperature", minimum=0.0, maximum=1.0, value=0.7, step=0.01),
-                             gr.Slider(label="Max new tokens", minimum=1, maximum=512, value=160),
-                             gr.Slider(label="Top-p", minimum=0.0, maximum=1.0, value=0.90),
-                             gr.Slider(label="Repetition Penalty", minimum=0.1, maximum=2.0, value=1.0)],
-                     outputs=chatbot_ui,
-                     live=True)
+additional_inputs=[
+    gr.Slider(
+        label="Temperature",
+        value=0.7,
+        minimum=0.0,
+        maximum=1.0,
+        step=0.01,
+        interactive=True,
+        info="Higher values generate more diverse outputs",
+    ),
+    gr.Slider(
+        label="Max new tokens",
+        value=6400,
+        minimum=0,
+        maximum=8000,
+        step=64,
+        interactive=True,
+        info="The maximum numbers of new tokens",
+    ),
+    gr.Slider(
+        label="Top-p",
+        value=0.90,
+        minimum=0.0,
+        maximum=1,
+        step=0.01,
+        interactive=True,
+        info="Higher values sample more low-probability tokens",
+    ),
+    gr.Slider(
+        label="Repetition penalty",
+        value=1.0,
+        minimum=0.1,
+        maximum=2.0,
+        step=0.1,
+        interactive=True,
+        info="Penalize repeated tokens",
+    )
+]
 
+iface = gr.ChatInterface(fn=generate, 
+                        chatbot=mychatbot,
+						additional_inputs=additional_inputs,
+                        retry_btn=None,
+                        undo_btn=None
+                       )
 
-with gr.Blocks() as app:
+with gr.Blocks() as demo:
     gr.HTML("<center><h1>Chat with GEMMA 7B</h1></center>")
     iface.render()
-
-app.launch()
+    
+demo.queue().launch(show_api=False)