|
import os |
|
import math |
|
|
|
import gradio as gr |
|
|
|
|
|
def get_chatbot_name(base_model, model_path_llama, inference_server='', debug=False): |
|
if not debug: |
|
inference_server = '' |
|
else: |
|
inference_server = ' : ' + inference_server |
|
if base_model == 'llama': |
|
model_path_llama = os.path.basename(model_path_llama) |
|
if model_path_llama.endswith('?download=true'): |
|
model_path_llama = model_path_llama.replace('?download=true', '') |
|
return f'h2oGPT [Model: {model_path_llama}{inference_server}]' |
|
else: |
|
return f'h2oGPT [Model: {base_model}{inference_server}]' |
|
|
|
|
|
def get_avatars(base_model, model_path_llama, inference_server=''): |
|
if base_model == 'llama': |
|
base_model = model_path_llama |
|
if inference_server is None: |
|
inference_server = '' |
|
|
|
model_base = os.getenv('H2OGPT_MODEL_BASE', 'models/') |
|
human_avatar = "human.jpg" |
|
if 'h2ogpt-gm'.lower() in base_model.lower(): |
|
bot_avatar = "h2oai.png" |
|
elif 'mistralai'.lower() in base_model.lower() or \ |
|
'mistral'.lower() in base_model.lower() or \ |
|
'mixtral'.lower() in base_model.lower(): |
|
bot_avatar = "mistralai.png" |
|
elif '01-ai/Yi-'.lower() in base_model.lower(): |
|
bot_avatar = "yi.svg" |
|
elif 'wizard' in base_model.lower(): |
|
bot_avatar = "wizard.jpg" |
|
elif 'openchat' in base_model.lower(): |
|
bot_avatar = "openchat.png" |
|
elif 'vicuna' in base_model.lower(): |
|
bot_avatar = "vicuna.jpeg" |
|
elif 'longalpaca' in base_model.lower(): |
|
bot_avatar = "longalpaca.png" |
|
elif 'llama2-70b-chat' in base_model.lower(): |
|
bot_avatar = "meta.png" |
|
elif 'llama2-13b-chat' in base_model.lower(): |
|
bot_avatar = "meta.png" |
|
elif 'llama2-7b-chat' in base_model.lower(): |
|
bot_avatar = "meta.png" |
|
elif 'llama2' in base_model.lower(): |
|
bot_avatar = "lama2.jpeg" |
|
elif 'llama-2' in base_model.lower(): |
|
bot_avatar = "lama2.jpeg" |
|
elif 'llama' in base_model.lower(): |
|
bot_avatar = "lama.jpeg" |
|
elif 'openai' in base_model.lower() or 'openai' in inference_server.lower(): |
|
bot_avatar = "openai.png" |
|
elif 'hugging' in base_model.lower(): |
|
bot_avatar = "hf-logo.png" |
|
elif 'claude' in base_model.lower(): |
|
bot_avatar = "anthropic.jpeg" |
|
elif 'gemini' in base_model.lower(): |
|
bot_avatar = "google.png" |
|
else: |
|
bot_avatar = "h2oai.png" |
|
|
|
bot_avatar = os.path.join(model_base, bot_avatar) |
|
human_avatar = os.path.join(model_base, human_avatar) |
|
|
|
human_avatar = human_avatar if os.path.isfile(human_avatar) else None |
|
bot_avatar = bot_avatar if os.path.isfile(bot_avatar) else None |
|
return human_avatar, bot_avatar |
|
|
|
|
|
def make_chatbots(output_label0, output_label0_model2, **kwargs): |
|
visible_models = kwargs['visible_models'] |
|
all_models = kwargs['all_possible_visible_models'] |
|
|
|
text_outputs = [] |
|
chat_kwargs = [] |
|
min_width = 250 if kwargs['gradio_size'] in ['small', 'large', 'medium'] else 160 |
|
for model_state_locki, model_state_lock in enumerate(kwargs['model_states']): |
|
output_label = get_chatbot_name(model_state_lock["base_model"], |
|
model_state_lock['llamacpp_dict']["model_path_llama"], |
|
model_state_lock["inference_server"], |
|
debug=bool(os.environ.get('DEBUG_MODEL_LOCK', 0))) |
|
if kwargs['avatars']: |
|
avatar_images = get_avatars(model_state_lock["base_model"], |
|
model_state_lock['llamacpp_dict']["model_path_llama"], |
|
model_state_lock["inference_server"]) |
|
else: |
|
avatar_images = None |
|
chat_kwargs.append(dict(render_markdown=kwargs.get('render_markdown', True), |
|
label=output_label, |
|
show_label=kwargs.get('visible_chatbot_label', True), |
|
elem_classes='chatsmall', |
|
height=kwargs['height'] or 400, |
|
min_width=min_width, |
|
avatar_images=avatar_images, |
|
show_copy_button=kwargs['show_copy_button'], |
|
visible=kwargs['model_lock'] and (visible_models is None or |
|
model_state_locki in visible_models or |
|
all_models[model_state_locki] in visible_models |
|
))) |
|
|
|
|
|
if visible_models and kwargs['model_lock_layout_based_upon_initial_visible']: |
|
len_visible = len(visible_models) |
|
else: |
|
len_visible = len(kwargs['model_states']) |
|
if kwargs['model_lock_columns'] == -1: |
|
kwargs['model_lock_columns'] = len_visible |
|
if kwargs['model_lock_columns'] is None: |
|
kwargs['model_lock_columns'] = 3 |
|
|
|
ncols = kwargs['model_lock_columns'] |
|
if kwargs['model_states'] == 0: |
|
nrows = 0 |
|
else: |
|
nrows = math.ceil(len_visible / kwargs['model_lock_columns']) |
|
|
|
if kwargs['model_lock_columns'] == 0: |
|
|
|
pass |
|
elif nrows <= 1: |
|
with gr.Row(): |
|
for chat_kwargs1, model_state_lock in zip(chat_kwargs, kwargs['model_states']): |
|
text_outputs.append(gr.Chatbot(**chat_kwargs1)) |
|
elif nrows == kwargs['model_states']: |
|
with gr.Row(): |
|
for chat_kwargs1, model_state_lock in zip(chat_kwargs, kwargs['model_states']): |
|
text_outputs.append(gr.Chatbot(**chat_kwargs1)) |
|
elif nrows > 0: |
|
len_chatbots = len(kwargs['model_states']) |
|
nrows = math.ceil(len_chatbots / kwargs['model_lock_columns']) |
|
for nrowi in range(nrows): |
|
with gr.Row(): |
|
for mii, (chat_kwargs1, model_state_lock) in enumerate(zip(chat_kwargs, kwargs['model_states'])): |
|
if mii < nrowi * len_chatbots / nrows or mii >= (1 + nrowi) * len_chatbots / nrows: |
|
continue |
|
text_outputs.append(gr.Chatbot(**chat_kwargs1)) |
|
if len(kwargs['model_states']) > 0: |
|
assert len(text_outputs) == len(kwargs['model_states']) |
|
|
|
if kwargs['avatars']: |
|
avatar_images = get_avatars(kwargs["base_model"], kwargs['llamacpp_dict']["model_path_llama"], |
|
kwargs["inference_server"]) |
|
else: |
|
avatar_images = None |
|
no_model_lock_chat_kwargs = dict(render_markdown=kwargs.get('render_markdown', True), |
|
show_label=kwargs.get('visible_chatbot_label', True), |
|
elem_classes='chatsmall', |
|
height=kwargs['height'] or 400, |
|
min_width=min_width, |
|
show_copy_button=kwargs['show_copy_button'], |
|
avatar_images=avatar_images, |
|
) |
|
with gr.Row(): |
|
text_output = gr.Chatbot(label=output_label0, |
|
visible=not kwargs['model_lock'], |
|
**no_model_lock_chat_kwargs, |
|
) |
|
text_output2 = gr.Chatbot(label=output_label0_model2, |
|
visible=False and not kwargs['model_lock'], |
|
**no_model_lock_chat_kwargs) |
|
return text_output, text_output2, text_outputs |
|
|
