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 tabs.arena_battle import get_tab_arena_battle
from tabs.arena_side_by_side import get_tab_arena_side_by_side
from tabs.playground import get_tab_playground

from models.chatml import ChatML
import json
import os

import gspread
from oauth2client.service_account import ServiceAccountCredentials

scope = ["https://spreadsheets.google.com/feeds", 'https://www.googleapis.com/auth/spreadsheets',
         "https://www.googleapis.com/auth/drive.file", "https://www.googleapis.com/auth/drive"]

GOOGLE_SHEETS_CERTIFICATE = json.loads(os.environ.get("GOOGLE_SHEETS_CERTIFICATE"))
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("\{\}", "{}")

creds = ServiceAccountCredentials.from_json_keyfile_dict(GOOGLE_SHEETS_CERTIFICATE, scope)
client = gspread.authorize(creds)

models = [
    BaseModel(
        name="PygmalionAI/pygmalion-13b",
        endpoint="pygmalion-13b",
        namespace="tenant-chaiml-guanaco",
        generation_params={
            'temperature': 0.7,
            'repetition_penalty': 1.0,
            'max_new_tokens': 128,
            'top_k': 10,
            'top_p': 0.9,
            'do_sample': True,
            'eos_token_id': 13,
        }
    ),
    BaseModel(
        name="lmsys/vicuna-7b-delta-v1.1",
        endpoint="vicuna-7b",
        namespace="tenant-chairesearch-test",
        generation_params={
            'temperature': 0.7,
            'repetition_penalty': 1.0,
            'max_new_tokens': 128,
            'top_k': 10,
            'top_p': 0.9,
            'do_sample': True,
            'eos_token_id': 13,
        }
    ),
    BaseModel(
        name="PygmalionAI/pygmalion-7b",
        endpoint="pygmalion-7b",
        namespace="tenant-chairesearch-test",
        generation_params={
            'temperature': 0.7,
            'repetition_penalty': 1.0,
            'max_new_tokens': 128,
            'top_k': 10,
            'top_p': 0.9,
            'do_sample': True,
            'eos_token_id': 13,
        }
    ),
    BaseModel(
        name="mosaicml/mpt-7b",
        endpoint="mpt-7b",
        namespace="tenant-chairesearch-test",
        generation_params={
            'temperature': 0.7,
            '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.7,
            '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.7,
            '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.7,
            '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.7,
            'repetition_penalty': 1.0,
            'max_new_tokens': 64,
            'top_k': 10,
            'top_p': 0.9,
            'do_sample': True,
            'eos_token_id': 187,
        }
    ),
]
model_mapping = {model.name: model for model in models}
print(list(model_mapping.keys()))


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:
    gr.Markdown("""
    # Chai: Model Evaluation
    Visit each tab for details ⬇️
    """)
    with gr.Tabs():
        with gr.TabItem("Playground"):
            get_tab_playground(download_bot_config, get_bot_profile, model_mapping)
        with gr.TabItem("Chatbot Arena (battle)"):
            get_tab_arena_battle(download_bot_config, get_bot_profile, model_mapping, client)
        with gr.TabItem("Chatbot Arena (side-by-side)"):
            get_tab_arena_side_by_side(download_bot_config, get_bot_profile, model_mapping, client)

demo.launch(enable_queue=False)