Update app.py

app.py

@@ -1,16 +1,23 @@
-from huggingface_hub import InferenceClient
 import gradio as gr
+from openai import OpenAI
+import os
 
 css = '''
-.gradio-container{max-width: 690px !important}
+.gradio-container{max-width: 1000px !important}
 h1{text-align:center}
 footer {
     visibility: hidden
 }
 '''
 
-client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3")
+ACCESS_TOKEN = os.getenv("HF_TOKEN")
 
+client = OpenAI(
+    base_url="https://api-inference.huggingface.co/v1/",
+    api_key=ACCESS_TOKEN,
+)
+
+# Mood prompts dictionary
 mood_prompts = {
     "Fun": "Respond in a light-hearted, playful manner.",
     "Serious": "Respond in a thoughtful, serious tone.",
@@ -62,74 +69,65 @@ mood_prompts = {
     "Worried": "Respond with concern and apprehension."
 }
 
-def format_prompt(message, history, system_prompt=None, mood=None):
-    prompt = "<s>"
-    if mood:
-        mood_description = mood_prompts.get(mood, "")
-        prompt += f"[SYS] {mood_description} [/SYS] "
-    for user_prompt, bot_response in history:
-        prompt += f"[INST] {user_prompt} [/INST]"
-        prompt += f" {bot_response}</s> "
-    if system_prompt:
-        prompt += f"[SYS] {system_prompt} [/SYS]"
-    prompt += f"[INST] {message} [/INST]"
-    return prompt
-
-def generate(
-    prompt, history, system_prompt=None, mood=None, temperature=0.2, max_new_tokens=1024, top_p=0.95, repetition_penalty=1.0,
+def respond(
+    message,
+    history: list[tuple[str, str]],
+    system_message,
+    max_tokens,
+    temperature,
+    top_p,
+    mood
 ):
-    temperature = float(temperature)
-    if temperature < 1e-2:
-        temperature = 1e-2
-    top_p = float(top_p)
-
-    generate_kwargs = dict(
-        temperature=temperature,
-        max_new_tokens=max_new_tokens,
-        top_p=top_p,
-        repetition_penalty=repetition_penalty,
-        do_sample=True,
-        seed=42,
-    )
-
-    formatted_prompt = format_prompt(prompt, history, system_prompt, mood)
-    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
-    output = ""
+    # Update system message with mood prompt
+    mood_prompt = mood_prompts.get(mood, "")
+    full_system_message = f"{system_message} {mood_prompt}".strip()
 
-    for response in stream:
-        output += response.token.text
-        yield output
+    messages = [{"role": "system", "content": full_system_message}]
 
-    # Append the latest interaction to history in tuple format
-    history.append((prompt, output))
-    return history  # Return updated history for output
+    for val in history:
+        if val[0]:
+            messages.append({"role": "user", "content": val[0]})
+        if val[1]:
+            messages.append({"role": "assistant", "content": val[1]})
 
-def gradio_interface():
-    with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
-        # Initialize state for history
-        history = gr.State([])
+    messages.append({"role": "user", "content": message})
 
-        # Row for system prompt and user prompt
-        with gr.Row():
-            system_prompt = gr.Textbox(placeholder="System prompt (optional)", lines=1, visible=False)
-            prompt = gr.Textbox(placeholder="Enter your message", lines=4)
+    response = ""
+    
+    for message in client.chat.completions.create(
+        model="meta-llama/Meta-Llama-3.1-8B-Instruct",
+        max_tokens=max_tokens,
+        stream=True,
+        temperature=temperature,
+        top_p=top_p,
+        messages=messages,
+    ):
+        token = message.choices[0].delta.content
         
-        # Row for generate button and output
-        with gr.Row():
-            generate_btn = gr.Button("Generate")
-            output = gr.Chatbot()
-
-        # Row for mood selection
-        with gr.Row():
-            mood = gr.Radio(choices=list(mood_prompts.keys()), value="Professional", label="Select Mood")
+        response += token
+        yield response
 
-        # Connect button click to generate function
-        generate_btn.click(
-            generate,
-            inputs=[prompt, history, system_prompt, mood],
-            outputs=[output]
+demo = gr.ChatInterface(
+    respond,
+    additional_inputs=[
+        gr.Textbox(value="", label="System message"),
+        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",
+        ),
+        gr.Dropdown(
+            choices=list(mood_prompts.keys()),
+            label="Mood",
+            value="Casual"
         )
-        
-    demo.queue().launch(show_api=False)
-
-gradio_interface()
+    ],
+    css=css,
+    theme="bethecloud/storj_theme",
+)
+if __name__ == "__main__":
+    demo.launch()