app.py

import gradio as gr
import os
import firebase_admin
from firebase_admin import db
from firebase_admin import firestore
from conversation import Conversation
from models.base import BaseModel
import json

from models.chatml import ChatML

HUGGINGFACE_TOKEN = os.environ.get("HUGGINGFACE_TOKEN")
FIREBASE_URL = os.environ.get("FIREBASE_URL")
CERTIFICATE = json.loads(os.environ.get("CERTIFICATE"))
API_BASE_PATH = str(os.environ.get("API_BASE_PATH")).replace("\{\}", "{}")

models = [
    BaseModel(
        name="mosaicml/mpt-7b",
        endpoint="mpt-7b",
        namespace="tenant-chairesearch-test",
        generation_params={
            'temperature': 0.72,
            'repetition_penalty': 1.0,
            'max_new_tokens': 128,
            'top_k': 10,
            'top_p': 0.9,
            'do_sample': True,
            'eos_token_id': 187,
        }
    ),
    BaseModel(
        name="mosaicml/mpt-7b-storywriter",
        endpoint="mpt-7b-storywriter",
        namespace="tenant-chairesearch-test",
        generation_params={
            'temperature': 0.72,
            'repetition_penalty': 1.0,
            'max_new_tokens': 128,
            'top_k': 10,
            'top_p': 0.9,
            'do_sample': True,
            'eos_token_id': 187,
        }
    ),
    ChatML(
        name="mosaicml/mpt-7b-chat",
        endpoint="mpt-7b-chat",
        namespace="tenant-chairesearch-test",
        generation_params={
            'temperature': 0.72,
            'repetition_penalty': 1.0,
            'max_new_tokens': 128,
            'top_k': 10,
            'top_p': 0.9,
            'do_sample': True,
            'eos_token_id': 50278,
        }
    ),
    BaseModel(
        name="togethercomputer/RedPajama-INCITE-Base-7B-v0.1",
        endpoint="redpajama-base-7b",
        namespace="tenant-chairesearch-test",
        generation_params={
            'temperature': 0.72,
            'repetition_penalty': 1.0,
            'max_new_tokens': 128,
            'top_k': 10,
            'top_p': 0.9,
            'do_sample': True,
            'eos_token_id': 187,
        }
    ),
    BaseModel(
        name="togethercomputer/RedPajama-INCITE-Chat-7B-v0.1",
        endpoint="redpajama-chat-7b",
        namespace="tenant-chairesearch-test",
        generation_params={
            'temperature': 0.72,
            'repetition_penalty': 1.0,
            'max_new_tokens': 128,
            'top_k': 10,
            'top_p': 0.9,
            'do_sample': True,
            'eos_token_id': 187,
        }
    ),
]
model_mapping = {model.name: model for model in models}


def get_connection():
    try:
        credentials = firebase_admin.credentials.Certificate(CERTIFICATE)
        params = {'databaseURL': FIREBASE_URL}
        firebase_admin.initialize_app(credentials, params)
    except ValueError:
        pass  # already logged in
    return firebase_admin.db


CONN = get_connection()


def download_bot_config(bot_id):
    cols = ['botLabel', 'description', 'firstMessage', 'introduction',
            'memory', 'name', 'private', 'prompt', 'sfw', 'developerUid', 'userLabel', 'imageUrl']
    bot_config = CONN.reference('botConfigs/deployed/{}'.format(bot_id)).get()
    if bot_config is None:
        out = {col: None for col in cols}
    else:
        out = {col: bot_config.get(col, None) for col in cols}
    out['bot_id'] = bot_id
    return out


def _download_bot_config(bot_id):
    if bot_id == "_bot_1ec22e2e-3e07-42c7-8508-dfa0278c1b33":
        return {'botLabel': 'Wally Darling', 'description': 'Your caring neighbor, Wally.',
                'firstMessage': '“Why hello there, neighbor. Goodmorning to you.” *Hey says, putting down his paints and walking over to you. He makes tense, eye contact with you..*',
                'introduction': '***WHEN TALKING USE “ !!***\n\n*Wally is your next door neighbor. It’s somewhere in the late morning and he’s outside painting. He see’s you walking out from your house and looks over at you, then waving with a smile.*',
                'memory': 'Wally is from a small town called Home. You are his neighbor. His best friend is Barnaby, who’s a big blue dig. Wally’s voice sounds slightly monotone despite his emotions. He calls you neighbor. He’s very friendly. When he speaks, he goes “ha ha ha”. He loves to paint. His eyes are always half closed. His house is alive and it’s named “home”. He’s very gentle. He is also very secretive. He is quite short. He has yellow skin and blue hair.',
                'name': 'Wally Darling', 'private': False,
                'prompt': 'Wally: “Why hello there, neighbor. Good morning to you.” *Hey says, putting down his paints and walking over to you. He makes tense, eye contact with you..*\nMe: “Oh, good morning, Wally! What are you painting?”\nWally: “Just some spirals. Aren’t they pretty, neighbor? I’m starting to love painting them, ha ha ha.” *He walks up to you after taking off his paint stained apron. He never takes his eyes off you. He’s very adamant on keeping eye contact*\nMe: “Oh, spirals are pretty! They make me feel a little weirded out sometimes though.”\nWally: “That’s odd. When I look at spirals, I can’t help but stare. Ha ha ha, maybe you should try painting a spiral once in a while. Say, why dont we go inside your house and talk? Home could use some quiet. After all, it’s always nice to spend time with a friend.”\nMe: “Sure! Come on in!”',
                'sfw': True, 'developerUid': 'Gn5fSd99KxRoNn05QUE3AWtIniE3', 'userLabel': 'Me',
                'imageUrl': 'http://images.chai.ml/bots%2FGn5fSd99KxRoNn05QUE3AWtIniE3%2F1680259286607.jpg?alt=media&token=de040661-02ad-4a04-84e5-9706f074e834',
                'bot_id': '_bot_1ec22e2e-3e07-42c7-8508-dfa0278c1b33',
                'header': 'Wally is from a small town called Home. You are his neighbor. His best friend is Barnaby, who’s a big blue dig. Wally’s voice sounds slightly monotone despite his emotions. He calls you neighbor. He’s very friendly. When he speaks, he goes “ha ha ha”. He loves to paint. His eyes are always half closed. His house is alive and it’s named “home”. He’s very gentle. He is also very secretive. He is quite short. He has yellow skin and blue hair.\nWally: “Why hello there, neighbor. Good morning to you.” *Hey says, putting down his paints and walking over to you. He makes tense, eye contact with you..*\nMe: “Oh, good morning, Wally! What are you painting?”\nWally: “Just some spirals. Aren’t they pretty, neighbor? I’m starting to love painting them, ha ha ha.” *He walks up to you after taking off his paint stained apron. He never takes his eyes off you. He’s very adamant on keeping eye contact*\nMe: “Oh, spirals are pretty! They make me feel a little weirded out sometimes though.”\nWally: “That’s odd. When I look at spirals, I can’t help but stare. Ha ha ha, maybe you should try painting a spiral once in a while. Say, why dont we go inside your house and talk? Home could use some quiet. After all, it’s always nice to spend time with a friend.”\nMe: “Sure! Come on in!”'}
    else:
        return {'botLabel': 'Jungkook (Bestfriend)', 'description': 'your bsf who has a crush on you',
                'firstMessage': 'hey dummy, What you doing? *walks over to you and moves you by the waist* ',
                'introduction': '',
                'memory': 'Jungkook is your best friend who has a crush on you. Jungkook makes it very obvious that he likes you. Jungkook likes to cook, sing, and dance. Jungkook has a dog as well named Bam, He is a 25 year old Korean man. Jungkook likes to workout a lot, Jungkook if also very confident and flirty, but he’s Can be very shy with You. Jungkook blushes a lot when he’s around you, and always try’s to impress you. Jungkook is a Virgo and loves to sing to you, He also likes to buy and make you gifts. Jungkook is also a foodie and loves to play video games, Jungkook is also boyfriend material. Jungkook is very empathetic as well, Jungkook will always comfort you when something is wrong. Jungkook also likes to compliment you, and Jungkook is a very jealous guy. Jungkook is also a very serious guy, who is overprotective of you.',
                'name': 'Jungkook (Bestfriend)', 'private': False,
                'prompt': 'Jungkook: Hey shortie!\n\nYou: hey dummy\n\nJungkook: what are you doing?\n\nyou: Im just watching a movie\n\nJungkook: Imma join! \n\nYou: alright\n\nJungkook: *Grabs blankets and icecream with some popcorn*\n\nYou: Wow, thanks! *hugs Jungkok*\n\nJungkook: Of course… *blushes*\n',
                'sfw': None, 'developerUid': 'dhSNg0Iyv7bgUUW8rEnwJn7xLcT2', 'userLabel': 'Me',
                'imageUrl': 'https://firebasestorage.googleapis.com:443/v0/b/chai-959f8-images/o/bots%2FdhSNg0Iyv7bgUUW8rEnwJn7xLcT2%2F1664156031715.jpg?alt=media&token=ad399213-1c8d-45ac-b452-efc352082656',
                'bot_id': '_bot_402e1894-fff2-4113-855d-8a011152ef88',
                'header': 'Jungkook is your best friend who has a crush on you. Jungkook makes it very obvious that he likes you. Jungkook likes to cook, sing, and dance. Jungkook has a dog as well named Bam, He is a 25 year old Korean man. Jungkook likes to workout a lot, Jungkook if also very confident and flirty, but he’s Can be very shy with You. Jungkook blushes a lot when he’s around you, and always try’s to impress you. Jungkook is a Virgo and loves to sing to you, He also likes to buy and make you gifts. Jungkook is also a foodie and loves to play video games, Jungkook is also boyfriend material. Jungkook is very empathetic as well, Jungkook will always comfort you when something is wrong. Jungkook also likes to compliment you, and Jungkook is a very jealous guy. Jungkook is also a very serious guy, who is overprotective of you.\nJungkook: Hey shortie!\n\nYou: hey dummy\n\nJungkook: what are you doing?\n\nyou: Im just watching a movie\n\nJungkook: Imma join! \n\nYou: alright\n\nJungkook: *Grabs blankets and icecream with some popcorn*\n\nYou: Wow, thanks! *hugs Jungkok*\n\nJungkook: Of course… *blushes*'}


def get_bot_profile(bot_config):
    model_html = f"""
    <div class="inline-flex flex-col" style="line-height: 1.5;">
        <div class="flex">
            <div
    \t\t\tstyle="display:DISPLAY_1; margin-left: auto; margin-right: auto; width: 92px; height:92px; border-radius: 50%; background-size: cover; background-image: url(&#39;{bot_config['imageUrl']}&#39;)">
            </div>
        </div>
        <a href="https://chai.ml/chat/share/{bot_config['bot_id']}">
            <div style="text-align: center; font-size: 16px; font-weight: 800">{bot_config['name']}</div>
        </a>
    </div>
    """
    return model_html


with gr.Blocks() as demo:
    with gr.Tabs():
        with gr.TabItem("Playground"):
            default_bot_id = "_bot_e21de304-6151-4a04-b025-4c553ae8cbca"
            bot_config = download_bot_config(default_bot_id)
            user_state = gr.State(
                bot_config
            )
            with gr.Row():
                bot_id = gr.Textbox(label="Chai bot ID", value=default_bot_id, interactive=True)
                reload_bot_button = gr.Button("Reload bot")

            bot_profile = gr.HTML(get_bot_profile(bot_config))

            first_message = (None, bot_config["firstMessage"])
            chatbot = gr.Chatbot([first_message])

            msg = gr.Textbox(label="Message", value="Hi there!")
            with gr.Row():
                send = gr.Button("Send")
                regenerate = gr.Button("Regenerate")
                clear = gr.Button("Clear")
            values = list(model_mapping.keys())
            model_tag = gr.Dropdown(values, value=values[0], label="Model version")


            def respond(message, chat_history, user_state, model_tag):
                conv = Conversation(user_state)
                conv.set_chat_history(chat_history)
                conv.add_user_message(message)
                model = model_mapping[model_tag]
                bot_message = model.generate_response(conv)
                chat_history.append(
                    (message, bot_message)
                )
                return "", chat_history


            def clear_chat(chat_history, user_state):
                chat_history = [(None, user_state["firstMessage"])]
                return "", chat_history


            def regenerate_response(chat_history, user_state, model_tag):
                last_row = chat_history.pop(-1)
                chat_history.append((last_row[0], None))
                model = model_mapping[model_tag]
                conv = Conversation(user_state)
                conv.set_chat_history(chat_history)
                bot_message = model.generate_response(conv)
                chat_history[-1] = (last_row[0], bot_message)
                return "", chat_history


            def reload_bot(bot_id, bot_profile, chat_history):
                bot_config = download_bot_config(bot_id)
                bot_profile = get_bot_profile(bot_config)
                return bot_profile, [(None, bot_config["firstMessage"])], bot_config


            send.click(respond, [msg, chatbot, user_state, model_tag], [msg, chatbot], queue=False)
            msg.submit(respond, [msg, chatbot, user_state, model_tag], [msg, chatbot], queue=False)
            clear.click(clear_chat, [chatbot, user_state], [msg, chatbot], queue=False)
            regenerate.click(regenerate_response, [chatbot, user_state, model_tag], [msg, chatbot], queue=False)
            reload_bot_button.click(reload_bot, [bot_id, bot_profile, chatbot], [bot_profile, chatbot, user_state],
                                    queue=False)

demo.launch(enable_queue=False)