Update app.py

app.py

@@ -28,9 +28,9 @@ def rewrite_csv_ordered_by_winning_rate(csv_path):
     # Save the sorted DataFrame to a new CSV file
     df_sorted.to_csv(csv_path, index=False)
 
-@spaces.GPU()
-def run_inference(pipe, prompt):
-    response = pipe(prompt)
+@spaces.GPU(duration=200)
+def run_inference(pipeline, prompt):
+    response = pipeline(prompt)
     bot_message = response[0]["generated_text"]
     return bot_message
 
@@ -39,7 +39,7 @@ def modelA_button():
     df = pd.read_csv("models.csv")
     df.loc[df["MODEL"] == choice["ModelA"], "MATCHES_WON"] += 1
     df.loc[df["MODEL"] == choice["ModelA"], "WINNING_RATE"] = df.loc[df["MODEL"] == choice["ModelA"], "MATCHES_WON"]/df.loc[df["MODEL"] == choice["ModelA"], "MATCHES_PLAYED"]  
-    df.to_csv("models.csv", index=False)
+    df.to_csv("models.csv")
     rewrite_csv_ordered_by_winning_rate("models.csv")
 
 def modelB_button():
@@ -47,34 +47,41 @@ def modelB_button():
     df = pd.read_csv("models.csv")
     df.loc[df["MODEL"] == choice["ModelB"], "MATCHES_WON"] += 1
     df.loc[df["MODEL"] == choice["ModelB"], "WINNING_RATE"] = df.loc[df["MODEL"] == choice["ModelB"], "MATCHES_WON"]/df.loc[df["MODEL"] == choice["ModelB"], "MATCHES_PLAYED"]  
-    df.to_csv("models.csv", index=False)
+    df.to_csv("models.csv")
     rewrite_csv_ordered_by_winning_rate("models.csv")
       
-def reply(modelA, modelB, prompt):
+import time
+
+def replyA(prompt, history, modelA):
     global choice
     choice["ModelA"] = modelA
-    choice["ModelB"] = modelB
     df = pd.read_csv("models.csv")
     df.loc[df["MODEL"] == modelA, "MATCHES_PLAYED"] += 1
-    df.loc[df["MODEL"] == modelB, "MATCHES_PLAYED"] += 1
     df.to_csv("models.csv", index=False)
-    pipeA = pipeline("text-generation", model=models_and_tokenizers[modelA][0], tokenizer=models_and_tokenizers[modelA][1], max_new_tokens=512, repetition_penalty=1.5, temperature=0.5, device_map="cuda:0")
+    pipeA = pipeline("text-generation", model=models_and_tokenizers[modelA][0], tokenizer=models_and_tokenizers[modelA][1], max_new_tokens=512, repetition_penalty=1.5, temperature=0.5, device="cuda")
     responseA = run_inference(pipeA, prompt)    
-    pipeB = pipeline("text-generation", model=models_and_tokenizers[modelB][0], tokenizer=models_and_tokenizers[modelB][1], max_new_tokens=512, repetition_penalty=1.5, temperature=0.5, device_map="cuda:1")
-    responseB = run_inference(pipeB, prompt)    
-    return responseA, responseB
+    r = ''
+    for c in responseA:
+        r+=c
+        time.sleep(0.0001)
+        yield r
 
-modelA_dropdown = gr.Dropdown(models_checkpoints, label="Model A", info="Choose the first model for the battle!")
-modelB_dropdown = gr.Dropdown(models_checkpoints, label="Model B", info="Choose the second model for the battle!")
-prompt_textbox = gr.Textbox(label="Prompt", value="Is pineapple pizza sacrilegious?")
+def replyB(prompt, history, modelB):
+    global choice
+    choice["ModelB"] = modelB
+    df = pd.read_csv("models.csv")
+    df.loc[df["MODEL"] == modelB, "MATCHES_PLAYED"] += 1
+    df.to_csv("models.csv", index=False)
+    pipeB = pipeline("text-generation", model=models_and_tokenizers[modelB][0], tokenizer=models_and_tokenizers[modelB][1], max_new_tokens=512, repetition_penalty=1.5, temperature=0.5, device="cuda")
+    responseB = run_inference(pipeB, prompt)    
+    r = ''
+    for c in responseB:
+        r+=c
+        time.sleep(0.0001)
+        yield r
 
-
-with gr.Blocks() as demo1:
-    demo0 = gr.Interface(fn=reply, inputs=[modelA_dropdown, modelB_dropdown, prompt_textbox], outputs=[gr.Textbox(label="Model A response"), gr.Textbox(label="Model B response")])
-    btnA = gr.Button("Vote for Model A!") 
-    btnB = gr.Button("Vote for Model B!")
-    btnA.click(modelA_button, inputs=None, outputs=None) 
-    btnB.click(modelB_button, inputs=None, outputs=None) 
+modelAchoice = gr.Dropdown(models_checkpoints, label="Model A")
+modelBchoice = gr.Dropdown(models_checkpoints, label="Model B")
 
 with gr.Blocks() as demo2:
     f = open("tab.html")
@@ -84,7 +91,24 @@ with gr.Blocks() as demo2:
     btn = gr.Button("Refresh")
     btn.click(fn=refreshfn, inputs=None, outputs=t)
 
-demo = gr.TabbedInterface([demo1, demo2], ["Chat Arena", "Leaderboard"])
+
+accrdnA = gr.Accordion(label="Choose model A")
+accrdnB = gr.Accordion(label="Choose model B")
+
+chtbA = gr.Chatbot(label="Chat with Model A")
+chtbB = gr.Chatbot(label="Chat with Model B")
+
+with gr.Blocks() as demo1:
+    gr.ChatInterface(fn=replyA, chatbot=chtbA, additional_inputs=modelAchoice, additional_inputs_accordion=accrdnA, submit_btn="Submit to Model A")
+    gr.ChatInterface(fn=replyB, chatbot=chtbB, additional_inputs=modelBchoice, additional_inputs_accordion=accrdnB, submit_btn="Submit to Model B")
+    btnA = gr.Button("Vote for Model A!") 
+    btnB = gr.Button("Vote for Model B!")
+    btnA.click(modelA_button, inputs=None, outputs=None) 
+    btnB.click(modelB_button, inputs=None, outputs=None) 
+
+demo = gr.TabbedInterface([demo1, demo2], ["Chat Arena", "Leaderboard"], title="""<h1 align='center'>SmolLM Arena</h1>
+<h2 align='center'>Cast your vote to choose the best Small Language Model (100M-1.7B)!🚀</h2>
+<h3 align='center'>[<a href="https://github.com/AstraBert/smollm-arena">GitHub</a>] [<a href="https://github.com/AstraBert/smollm-arena?tab=readme-ov-file#usage">Usage Guide</a>]""")
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)