app.py

import gradio as gr
import random
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("models/" + new_models[inp[z]]))
        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]
        client = gr.load()
        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]
        client1=gr.load(models[0])
        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]
        client2=gr.load(models[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]
        client3=gr.load(models[2])
        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]
        client4=gr.load(models[3])
        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()