Refactor model ID handling in app.py and update requirements.txt

app.py

@@ -2,6 +2,7 @@ import spaces
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import torch
 import gradio as gr
+
 from threading import Thread
 import subprocess
 subprocess.run('pip install -U flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
@@ -11,9 +12,9 @@ models_available = [
     "silma-ai/SILMA-9B-Instruct-v1.0",
     "inceptionai/jais-adapted-7b-chat",
     "inceptionai/jais-family-6p7b-chat",
-    "NousResearch/Meta-Llama-3.1-8B-Instruct",
-    # "unsloth/gemma-2-9b-it",
-    "NousResearch/Meta-Llama-3-8B-Instruct",
+    "inceptionai/jais-family-2p7b-chat",
+    "inceptionai/jais-family-1p3b-chat",
+    "inceptionai/jais-family-590m-chat",
 ]
 
 tokenizer_a, model_a = None, None
@@ -69,6 +70,7 @@ def load_model_b(model_id):
 
 @spaces.GPU()
 def generate_both(system_prompt, input_text, chatbot_a, chatbot_b, max_new_tokens=2048, temperature=0.2, top_p=0.9, repetition_penalty=1.1):
+
     text_streamer_a = TextIteratorStreamer(tokenizer_a, skip_prompt=True)
     text_streamer_b = TextIteratorStreamer(tokenizer_b, skip_prompt=True)
 
@@ -163,7 +165,7 @@ def generate_both(system_prompt, input_text, chatbot_a, chatbot_b, max_new_token
 def clear():
     return [], []
 
-arena_notes = """Important Notes:
+arena_notes = """## Important Notes:
 - Sometimes an error may occur when generating the response, in this case, please try again.
 """
 
@@ -173,11 +175,12 @@ with gr.Blocks() as demo:
         gr.Markdown(arena_notes)
         system_prompt = gr.Textbox(lines=1, label="System Prompt", value="أنت متحدث لبق باللغة العربية!", rtl=True, text_align="right", show_copy_button=True)
         with gr.Row(variant="panel"):
-            model_dropdown_a = gr.Dropdown(label="Model A", choices=models_available, value=None)
-            model_dropdown_b = gr.Dropdown(label="Model B", choices=models_available, value=None)
-        with gr.Row(variant="panel"):
-            chatbot_a = gr.Chatbot(label="Model A", rtl=True, likeable=True, show_copy_button=True, height=500)
-            chatbot_b = gr.Chatbot(label="Model B", rtl=True, likeable=True, show_copy_button=True, height=500)
+            with gr.Column():
+                model_dropdown_a = gr.Dropdown(label="Model A", choices=models_available, value=None)
+                chatbot_a = gr.Chatbot(label="Model A", rtl=True, likeable=True, show_copy_button=True, height=500)
+            with gr.Column():
+                model_dropdown_b = gr.Dropdown(label="Model B", choices=models_available, value=None)
+                chatbot_b = gr.Chatbot(label="Model B", rtl=True, likeable=True, show_copy_button=True, height=500)
         with gr.Row(variant="panel"):
             with gr.Column(scale=1):
                 submit_btn = gr.Button(value="Generate", variant="primary")