Create app.py

app.py

@@ -0,0 +1,283 @@
+import gradio as gr
+
+
+models=[
+    "facebook/MobileLLM-125M",
+    "facebook/MobileLLM-350M",
+    "facebook/MobileLLM-600M",
+    "facebook/MobileLLM-1B",
+]
+client_z=[]
+
+
+def load_models(inp,new_models):
+    if not new_models:
+        new_models=models
+    out_box=[gr.Chatbot(),gr.Chatbot(),gr.Chatbot(),gr.Chatbot()]
+    print(type(inp))
+    print(inp)
+    #print(new_models[inp[0]])
+    client_z.clear()
+    for z,ea in enumerate(inp):
+        client_z.append(gr.load(new_models[inp[z]],token=token))
+        out_box[z]=(gr.update(label=new_models[inp[z]]))
+    return out_box[0],out_box[1],out_box[2],out_box[3]
+def format_prompt_default(message, history):
+    prompt = ""
+    if history:
+        #<start_of_turn>userHow does the brain work?<end_of_turn><start_of_turn>model
+        for user_prompt, bot_response in history:
+            prompt += f"{user_prompt}\n"
+            print(prompt)
+            prompt += f"{bot_response}\n"
+            print(prompt)
+    prompt += f"{message}\n"
+    return prompt
+
+def format_prompt_gemma(message, history):
+    prompt = ""
+    if history:
+        #<start_of_turn>userHow does the brain work?<end_of_turn><start_of_turn>model
+        for user_prompt, bot_response in history:
+            prompt += f"{user_prompt}\n"
+            print(prompt)
+            prompt += f"{bot_response}\n"
+            print(prompt)
+    prompt += f"<start_of_turn>user{message}<end_of_turn><start_of_turn>model"
+    return prompt
+
+   
+def format_prompt_mixtral(message, history):
+    prompt = "<s>"
+    if history:
+        for user_prompt, bot_response in history:
+            prompt += f"[INST] {user_prompt} [/INST]"
+            prompt += f" {bot_response}</s> "
+    prompt += f"[INST] {message} [/INST]"
+    return prompt
+
+def format_prompt_choose(message, history, model_name, new_models=None):
+    if not new_models:
+        new_models=models
+    if "gemma" in new_models[model_name].lower() and "it" in new_models[model_name].lower():
+        return format_prompt_gemma(message,history)
+    if "mixtral" in new_models[model_name].lower():
+        return format_prompt_mixtral(message,history)
+    else:
+        return format_prompt_mixtral(message,history)
+
+
+
+mega_hist=[[],[],[],[]]
+def chat_inf_tree(system_prompt,prompt,history,client_choice,seed,temp,tokens,top_p,rep_p,hid_val):
+    if len(client_choice)>=hid_val:
+        client=client_z[int(hid_val)-1]
+        if history:
+            mega_hist[hid_val-1]=history
+            #history = []
+            hist_len=0
+        generate_kwargs = dict(
+            temperature=temp,
+            max_new_tokens=tokens,
+            top_p=top_p,
+            repetition_penalty=rep_p,
+            do_sample=True,
+            seed=seed,
+        )
+        #formatted_prompt=prompt   
+        formatted_prompt = format_prompt(f"{system_prompt}, {prompt}", mega_hist[hid_val-1])
+        stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
+        output = ""
+        for response in stream:
+            output += response.token.text
+            yield [(prompt,output)]
+        mega_hist[hid_val-1].append((prompt,output))
+        yield mega_hist[hid_val-1]
+    else:
+        yield None
+
+
+
+
+def chat_inf_a(system_prompt,prompt,history,client_choice,seed,temp,tokens,top_p,rep_p,hid_val):
+    if len(client_choice)>=hid_val:
+        if system_prompt:
+            system_prompt=f'{system_prompt}, '
+        client1=client_z[int(hid_val)-1]
+        if not history:
+            history = []
+            hist_len=0
+        generate_kwargs = dict(
+            temperature=temp,
+            max_new_tokens=tokens,
+            top_p=top_p,
+            repetition_penalty=rep_p,
+            do_sample=True,
+            seed=seed,
+        )
+        #formatted_prompt=prompt   
+        formatted_prompt = format_prompt_choose(f"{system_prompt}{prompt}", history, client_choice[0])
+        stream1 = client1.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
+        output = ""
+        for response in stream1:
+            output += response.token.text
+            yield [(prompt,output)]
+        history.append((prompt,output))
+        yield history
+    else:
+        yield None
+
+
+def chat_inf_b(system_prompt,prompt,history,client_choice,seed,temp,tokens,top_p,rep_p,hid_val):
+    if len(client_choice)>=hid_val:
+        if system_prompt:
+            system_prompt=f'{system_prompt}, '        
+        client2=client_z[int(hid_val)-1]
+        if not history:
+            history = []
+            hist_len=0
+        generate_kwargs = dict(
+            temperature=temp,
+            max_new_tokens=tokens,
+            top_p=top_p,
+            repetition_penalty=rep_p,
+            do_sample=True,
+            seed=seed,
+        )
+        #formatted_prompt=prompt   
+        formatted_prompt = format_prompt_choose(f"{system_prompt}{prompt}", history, client_choice[1])
+        stream2 = client2.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
+        output = ""
+        for response in stream2:
+            output += response.token.text
+            yield [(prompt,output)]
+        history.append((prompt,output))
+        yield history
+    else:
+        yield None
+
+def chat_inf_c(system_prompt,prompt,history,client_choice,seed,temp,tokens,top_p,rep_p,hid_val):
+    if len(client_choice)>=hid_val:
+        if system_prompt:
+            system_prompt=f'{system_prompt}, '           
+        client3=client_z[int(hid_val)-1]
+        if not history:
+            history = []
+            hist_len=0
+        generate_kwargs = dict(
+            temperature=temp,
+            max_new_tokens=tokens,
+            top_p=top_p,
+            repetition_penalty=rep_p,
+            do_sample=True,
+            seed=seed,
+        )
+        #formatted_prompt=prompt   
+        formatted_prompt = format_prompt_choose(f"{system_prompt}{prompt}", history, client_choice[2])
+        stream3 = client3.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
+        output = ""
+        for response in stream3:
+            output += response.token.text
+            yield [(prompt,output)]
+        history.append((prompt,output))
+        yield history
+    else:
+        yield None
+
+def chat_inf_d(system_prompt,prompt,history,client_choice,seed,temp,tokens,top_p,rep_p,hid_val):
+    if len(client_choice)>=hid_val:
+        if system_prompt:
+            system_prompt=f'{system_prompt}, '           
+        client4=client_z[int(hid_val)-1]
+        if not history:
+            history = []
+            hist_len=0
+        generate_kwargs = dict(
+            temperature=temp,
+            max_new_tokens=tokens,
+            top_p=top_p,
+            repetition_penalty=rep_p,
+            do_sample=True,
+            seed=seed,
+        )
+        #formatted_prompt=prompt   
+        formatted_prompt = format_prompt_choose(f"{system_prompt}{prompt}", history, client_choice[3])
+        stream4 = client4.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
+        output = ""
+        for response in stream4:
+            output += response.token.text
+            yield [(prompt,output)]
+        history.append((prompt,output))
+        yield history
+    else:
+        yield None        
+def add_new_model(inp, cur):
+    cur.append(inp)
+    return cur,gr.update(choices=[z for z in cur])
+def load_new(models=models):
+    return models
+    
+def clear_fn():
+    return None,None,None,None,None,None
+rand_val=random.randint(1,1111111111111111)
+def check_rand(inp,val):
+    if inp==True:
+        return gr.Slider(label="Seed", minimum=1, maximum=1111111111111111, value=random.randint(1,1111111111111111))
+    else:
+        return gr.Slider(label="Seed", minimum=1, maximum=1111111111111111, value=int(val))
+
+with gr.Blocks() as app:
+    new_models=gr.State([])
+    gr.HTML("""<center><h1 style='font-size:xx-large;'>Chatbot Model Compare</h1>""")
+    with gr.Row():
+        chat_a = gr.Chatbot(height=500)
+        chat_b = gr.Chatbot(height=500)
+    with gr.Row():
+        chat_c = gr.Chatbot(height=500)
+        chat_d = gr.Chatbot(height=500)
+    with gr.Group():
+        with gr.Row():
+            with gr.Column(scale=3):
+                inp = gr.Textbox(label="Prompt")
+                sys_inp = gr.Textbox(label="System Prompt (optional)")
+                with gr.Row():
+                    with gr.Column(scale=2):
+                        btn = gr.Button("Chat")
+                    with gr.Column(scale=1):
+                        with gr.Group():
+                            stop_btn=gr.Button("Stop")
+                            clear_btn=gr.Button("Clear")                
+                client_choice=gr.Dropdown(label="Models",type='index', choices=[c for c in models],max_choices=4,multiselect=True,interactive=True)
+                add_model=gr.Textbox(label="New Model")
+                add_btn=gr.Button("Add Model")
+            with gr.Column(scale=1):
+                with gr.Group():
+                    rand = gr.Checkbox(label="Random Seed", value=True)
+                    seed=gr.Slider(label="Seed", minimum=1, maximum=1111111111111111,step=1, value=rand_val)
+                    tokens = gr.Slider(label="Max new tokens",value=3840,minimum=0,maximum=8000,step=64,interactive=True, visible=True,info="The maximum number of tokens")
+                    temp=gr.Slider(label="Temperature",step=0.01, minimum=0.01, maximum=1.0, value=0.9)
+                    top_p=gr.Slider(label="Top-P",step=0.01, minimum=0.01, maximum=1.0, value=0.9)
+                    rep_p=gr.Slider(label="Repetition Penalty",step=0.1, minimum=0.1, maximum=2.0, value=1.0)
+
+    hid1=gr.Number(value=1,visible=False)        
+    hid2=gr.Number(value=2,visible=False)        
+    hid3=gr.Number(value=3,visible=False)        
+    hid4=gr.Number(value=4,visible=False)    
+
+    app.load(load_new,None,new_models)
+    add_btn.click(add_new_model,[add_model,new_models],[new_models,client_choice])
+    client_choice.change(load_models,[client_choice,new_models],[chat_a,chat_b,chat_c,chat_d])
+
+    #im_go=im_btn.click(get_screenshot,[chat_b,im_height,im_width,chatblock,theme,wait_time],img)
+    #chat_sub=inp.submit(check_rand,[rand,seed],seed).then(chat_inf,[sys_inp,inp,chat_b,client_choice,seed,temp,tokens,top_p,rep_p],chat_b)
+    
+    go1=btn.click(check_rand,[rand,seed],seed).then(chat_inf_a,[sys_inp,inp,chat_b,client_choice,seed,temp,tokens,top_p,rep_p,hid1],chat_a)
+    go2=btn.click(check_rand,[rand,seed],seed).then(chat_inf_b,[sys_inp,inp,chat_b,client_choice,seed,temp,tokens,top_p,rep_p,hid2],chat_b)
+    go3=btn.click(check_rand,[rand,seed],seed).then(chat_inf_c,[sys_inp,inp,chat_b,client_choice,seed,temp,tokens,top_p,rep_p,hid3],chat_c)
+    go4=btn.click(check_rand,[rand,seed],seed).then(chat_inf_d,[sys_inp,inp,chat_b,client_choice,seed,temp,tokens,top_p,rep_p,hid4],chat_d)
+    
+    stop_btn.click(None,None,None,cancels=[go1,go2,go3,go4])
+    clear_btn.click(clear_fn,None,[inp,sys_inp,chat_a,chat_b,chat_c,chat_d])
+app.queue(default_concurrency_limit=10).launch()
+
+