Update app.py

app.py

@@ -58,7 +58,7 @@ def replyA(prompt, history, modelA):
     df = pd.read_csv("models.csv")
     df.loc[df["MODEL"] == modelA, "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="cuda")
+    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")
     responseA = run_inference(pipeA, prompt)    
     r = ''
     for c in responseA:
@@ -72,7 +72,7 @@ def replyB(prompt, history, 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")
+    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:0")
     responseB = run_inference(pipeB, prompt)    
     r = ''
     for c in responseB:
@@ -106,9 +106,7 @@ with gr.Blocks() as demo1:
     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>]""")
+demo = gr.TabbedInterface([demo1, demo2], ["Chat Arena", "Leaderboard"])
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)