feat: Add Arabic-ORPO-Llama3 chatbot comparison functionality

app.py

@@ -4,106 +4,162 @@ import torch
 import gradio as gr
 from threading import Thread
 
-base_model_id = "NousResearch/Meta-Llama-3-8B-Instruct"
-new_model_id = "MohamedRashad/Arabic-Orpo-Llama-3-8B-Instruct"
-
-# Reload tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained(base_model_id)
-base_model = AutoModelForCausalLM.from_pretrained(
-    base_model_id,
-    torch_dtype=torch.bfloat16,
-    device_map="auto",
-).eval()
-new_model = AutoModelForCausalLM.from_pretrained(
-    new_model_id,
-    torch_dtype=torch.bfloat16,
-    device_map="auto",
-).eval()
-terminators = [
-    tokenizer.eos_token_id,
-    tokenizer.convert_tokens_to_ids("<|eot_id|>"),
+models_available = [
+    "MohamedRashad/Arabic-Orpo-Llama-3-8B-Instruct",
+    "silma-ai/SILMA-9B-Instruct-v0.1.1",
+    "inceptionai/jais-adapted-7b-chat",
+    "inceptionai/jais-adapted-13b-chat",
+    "inceptionai/jais-family-6p7b-chat",
+    "inceptionai/jais-family-13b-chat",
 ]
 
+# Loop over all models to download them
+for model_id in models_available:
+    try:
+        AutoTokenizer.from_pretrained(model_id)
+        AutoModelForCausalLM.from_pretrained(model_id)
+    except:
+        models_available.remove(model_id)
+
+tokenizer_a, model_a = None, None
+tokenizer_b, model_b = None, None
+
+def load_model_a(model_id):
+    global tokenizer_a, model_a
+    tokenizer_a = AutoTokenizer.from_pretrained(model_id)
+    print(f"model A: {tokenizer_a.eos_token}")
+    try:
+        model_a = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            attn_implementation="flash_attention_2",
+            trust_remote_code=True,
+        ).eval()
+    except:
+        print(f"Using default attention implementation in {model_id}")
+        model_a = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            trust_remote_code=True,
+        ).eval()
+    return gr.update(label=model_id)
+    
+def load_model_b(model_id):
+    global tokenizer_b, model_b
+    tokenizer_b = AutoTokenizer.from_pretrained(model_id)
+    print(f"model B: {tokenizer_b.eos_token}")
+    try:
+        model_b = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            attn_implementation="flash_attention_2",
+            trust_remote_code=True,
+        ).eval()
+    except:
+        print(f"Using default attention implementation in {model_id}")
+        model_b = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            trust_remote_code=True,
+        ).eval()
+    return gr.update(label=model_id)
+
 @spaces.GPU(duration=120)
-def generate_both(system_prompt, input_text, base_chatbot, new_chatbot, max_new_tokens=2048, temperature=0.2, top_p=0.9, repetition_penalty=1.1):
-    base_text_streamer = TextIteratorStreamer(tokenizer, skip_prompt=True)
-    new_text_streamer = TextIteratorStreamer(tokenizer, skip_prompt=True)
+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)
 
     system_prompt_list = [{"role": "system", "content": system_prompt}]
     input_text_list = [{"role": "user", "content": input_text}]
 
-    base_chat_history = []
-    for user, assistant in base_chatbot:
-        base_chat_history.append({"role": "user", "content": user})
-        base_chat_history.append({"role": "assistant", "content": assistant})
+    chat_history_a = []
+    for user, assistant in chatbot_a:
+        chat_history_a.append({"role": "user", "content": user})
+        chat_history_a.append({"role": "assistant", "content": assistant})
 
-    new_chat_history = []
-    for user, assistant in new_chatbot:
-        new_chat_history.append({"role": "user", "content": user})
-        new_chat_history.append({"role": "assistant", "content": assistant})
+    chat_history_b = []
+    for user, assistant in chatbot_b:
+        chat_history_b.append({"role": "user", "content": user})
+        chat_history_b.append({"role": "assistant", "content": assistant})
     
-    base_messages = system_prompt_list + base_chat_history + input_text_list
-    new_messages = system_prompt_list + new_chat_history + input_text_list
+    base_messages = system_prompt_list + chat_history_a + input_text_list
+    new_messages = system_prompt_list + chat_history_b + input_text_list
 
-    base_input_ids = tokenizer.apply_chat_template(
+    input_ids_a = tokenizer_a.apply_chat_template(
         base_messages,
         add_generation_prompt=True,
         return_tensors="pt"
-    ).to(base_model.device).long()
+    ).to(model_a.device)
 
-    new_input_ids = tokenizer.apply_chat_template(
+    input_ids_b = tokenizer_b.apply_chat_template(
         new_messages,
         add_generation_prompt=True,
         return_tensors="pt"
-    ).to(new_model.device).long()
+    ).to(model_b.device)
 
-    base_generation_kwargs = dict(
-        input_ids=base_input_ids,
-        streamer=base_text_streamer,
+    generation_kwargs_a = dict(
+        input_ids=input_ids_a,
+        streamer=text_streamer_a,
         max_new_tokens=max_new_tokens,
-        eos_token_id=terminators,
-        pad_token_id=tokenizer.eos_token_id,
+        pad_token_id=tokenizer_a.eos_token_id,
         do_sample=True if temperature > 0 else False,
         temperature=temperature,
         top_p=top_p,
         repetition_penalty=repetition_penalty,
     )
-    new_generation_kwargs = dict(
-        input_ids=new_input_ids,
-        streamer=new_text_streamer,
+    generation_kwargs_b = dict(
+        input_ids=input_ids_b,
+        streamer=text_streamer_b,
         max_new_tokens=max_new_tokens,
-        eos_token_id=terminators,
-        pad_token_id=tokenizer.eos_token_id,
+        pad_token_id=tokenizer_b.eos_token_id,
         do_sample=True if temperature > 0 else False,
         temperature=temperature,
         top_p=top_p,
         repetition_penalty=repetition_penalty,
     )
 
-    base_thread = Thread(target=base_model.generate, kwargs=base_generation_kwargs)
-    base_thread.start()
-
-    base_chatbot.append([input_text, ""])
-    new_chatbot.append([input_text, ""])
-    
-    for base_text in base_text_streamer:
-        if "<|eot_id|>" in base_text:
-            eot_location = base_text.find("<|eot_id|>")
-            base_text = base_text[:eot_location]
-        base_chatbot[-1][-1] += base_text
-        yield base_chatbot, new_chatbot
-
-    new_thread = Thread(target=new_model.generate, kwargs=new_generation_kwargs)
-    new_thread.start()
-
-    for new_text in new_text_streamer:
-        if "<|eot_id|>" in new_text:
-            eot_location = new_text.find("<|eot_id|>")
-            new_text = new_text[:eot_location]
-        new_chatbot[-1][-1] += new_text
-        yield base_chatbot, new_chatbot
-
-    return base_chatbot, new_chatbot
+    thread_a = Thread(target=model_a.generate, kwargs=generation_kwargs_a)
+    thread_b = Thread(target=model_b.generate, kwargs=generation_kwargs_b)
+
+    thread_a.start()
+    thread_b.start()
+
+    chatbot_a.append([input_text, ""])
+    chatbot_b.append([input_text, ""])
+
+    finished_a = False
+    finished_b = False
+
+    while not (finished_a and finished_b):
+        if not finished_a:
+            try:
+                text_a = next(text_streamer_a)
+                if tokenizer_a.eos_token in text_a:
+                    eot_location = text_a.find(tokenizer_a.eos_token)
+                    text_a = text_a[:eot_location]
+                    finished_a = True
+                chatbot_a[-1][-1] += text_a
+                yield chatbot_a, chatbot_b
+            except StopIteration:
+                finished_a = True
+
+        if not finished_b:
+            try:
+                text_b = next(text_streamer_b)
+                if tokenizer_b.eos_token in text_b:
+                    eot_location = text_b.find(tokenizer_b.eos_token)
+                    text_b = text_b[:eot_location]
+                    finished_b = True
+                chatbot_b[-1][-1] += text_b
+                yield chatbot_a, chatbot_b
+            except StopIteration:
+                finished_b = True
+
+    return chatbot_a, chatbot_b
 
 def clear():
     return [], []
@@ -113,8 +169,11 @@ with gr.Blocks(title="Arabic-ORPO-Llama3") as demo:
         gr.HTML("<center><h1>Arabic Chatbot Comparison</h1></center>")
         system_prompt = gr.Textbox(lines=1, label="System Prompt", value="أنت متحدث لبق باللغة العربية!", rtl=True, text_align="right", show_copy_button=True)
         with gr.Row(variant="panel"):
-            base_chatbot = gr.Chatbot(label=base_model_id, rtl=True, likeable=True, show_copy_button=True, height=500)
-            new_chatbot = gr.Chatbot(label=new_model_id, rtl=True, likeable=True, show_copy_button=True, height=500)
+            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.Row(variant="panel"):
             with gr.Column(scale=1):
                 submit_btn = gr.Button(value="Generate", variant="primary")
@@ -126,8 +185,11 @@ with gr.Blocks(title="Arabic-ORPO-Llama3") as demo:
             top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.9, label="Top-p", step=0.01)
             repetition_penalty = gr.Slider(minimum=0.1, maximum=2.0, value=1.1, label="Repetition Penalty", step=0.1)
 
-    input_text.submit(generate_both, inputs=[system_prompt, input_text, base_chatbot, new_chatbot, max_new_tokens, temperature, top_p, repetition_penalty], outputs=[base_chatbot, new_chatbot])
-    submit_btn.click(generate_both, inputs=[system_prompt, input_text, base_chatbot, new_chatbot, max_new_tokens, temperature, top_p, repetition_penalty], outputs=[base_chatbot, new_chatbot])
-    clear_btn.click(clear, outputs=[base_chatbot, new_chatbot])
+    model_dropdown_a.change(load_model_a, inputs=[model_dropdown_a], outputs=[chatbot_a])
+    model_dropdown_b.change(load_model_b, inputs=[model_dropdown_b], outputs=[chatbot_b])
+
+    input_text.submit(generate_both, inputs=[system_prompt, input_text, chatbot_a, chatbot_b, max_new_tokens, temperature, top_p, repetition_penalty], outputs=[chatbot_a, chatbot_b])
+    submit_btn.click(generate_both, inputs=[system_prompt, input_text, chatbot_a, chatbot_b, max_new_tokens, temperature, top_p, repetition_penalty], outputs=[chatbot_a, chatbot_b])
+    clear_btn.click(clear, outputs=[chatbot_a, chatbot_b])
 
-demo.launch()
+demo.launch()