Update app.py

app.py

@@ -1,4 +1,4 @@
-import os
+
 import json
 import subprocess
 from threading import Thread
@@ -10,7 +10,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
-MODEL_ID = "nikravan/Marco_o1_q4"
+MODEL_ID = "nikravan/Marco-O1-q4"
 CHAT_TEMPLATE = "ChatML"
 MODEL_NAME = MODEL_ID.split("/")[-1]
 CONTEXT_LENGTH = 16000
@@ -120,26 +120,22 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 
 # Create Gradio interface
-with gr.Blocks(theme=gr.themes.Soft(primary_hue=COLOR)) as demo:
-    chatbot = gr.Chatbot(label=EMOJI + " " + MODEL_NAME, latex_delimiters=latex_delimiters_set)
-    system_prompt = gr.Textbox("You are a code assistant.", label="System prompt")
-    temperature = gr.Slider(0, 1, 0.3, label="Temperature")
-    max_new_tokens = gr.Slider(128, 4096, 1024, label="Max new tokens")
-    top_k = gr.Slider(1, 80, 40, label="Top K sampling")
-    repetition_penalty = gr.Slider(0, 2, 1.1, label="Repetition penalty")
-    top_p = gr.Slider(0, 1, 0.95, label="Top P sampling")
-    message = gr.Textbox(label="User Input")
-    submit = gr.Button("Submit")
-    
-    def respond(message, chatbot_history, system_prompt, temperature, max_new_tokens, top_k, repetition_penalty, top_p):
-        response = predict(message, chatbot_history, system_prompt, temperature, max_new_tokens, top_k, repetition_penalty, top_p)
-        chatbot_history.append((message, "".join(response)))
-        return chatbot_history
+gr.ChatInterface(
+    predict,
+    title=EMOJI + " " + MODEL_NAME,
+    description=DESCRIPTION,
     
-    submit.click(
-        respond, 
-        inputs=[message, chatbot, system_prompt, temperature, max_new_tokens, top_k, repetition_penalty, top_p],
-        outputs=chatbot
-    )
 
-demo.launch()
+     
+    additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=False),
+    additional_inputs=[
+        gr.Textbox("You are a code assistant.", label="System prompt"),
+        gr.Slider(0, 1, 0.3, label="Temperature"),
+        gr.Slider(128, 4096, 1024, label="Max new tokens"),
+        gr.Slider(1, 80, 40, label="Top K sampling"),
+        gr.Slider(0, 2, 1.1, label="Repetition penalty"),
+        gr.Slider(0, 1, 0.95, label="Top P sampling"),
+    ],
+    theme=gr.themes.Soft(primary_hue=COLOR),
+).queue().launch()
+