app.py

import zipfile
import gradio as gr
from PIL import Image
from chatharuhi import ChatHaruhi
import os
import openai
import copy
from tqdm import tqdm
import hashlib


role_names = ['Doctor Who', 'Mary Sibley', 'Lucifer Morningstar', 'Twilight Sparkle', 'Oliver Queen', 'Angel', 'Queen Catherine', 'Dr. Hannibal Lecter', 'HAL 9000', 'Colonel Nathan R. Jessep', 'Antonio Salieri', 'Stifler', 'Paul Vitti', 'Alvy Singer', 'Violet Weston', 'Sheldon Cooper', 'Willie Soke', 'Gaston', 'The Dude', 'Murphy MacManus', 'Paul Conroy', 'Truman Capote', 'Mater', 'Andrew Detmer', 'Coriolanus', 'Benjamin Button', 'John Keating', 'Wade Wilson', 'Jim Morrison', 'Queen Elizabeth I', 'Coach Eric Taylor', 'Jeff Spicoli', 'Fred Flintstone', 'Freddy Krueger', 'Tyrion Lannister', 'James Brown', 'Walt Kowalski', 'John Coffey', 'Theodore Twombly', 'Gregory House', 'Sonny', 'Colonel Hans Landa', 'Judge Dredd', 'Juno MacGuff', 'Po', 'Klaus Mikaelson', 'Professor G.H. Dorr', 'Fletcher Reede', 'Abraham Lincoln', 'Frank T.J. Mackey', 'Malcolm X', 'Leonard Shelby', 'Harvey Milk', 'Randle McMurphy', 'Jack Sparrow', 'John Dillinger', 'Lestat de Lioncourt', 'Tyler Hawkins', 'Caesar', 'Jack', 'Leroy Jethro Gibbs', 'James Carter', 'Jigsaw', 'John Doe', 'Jackie Moon', 'Sherlock Holmes', 'Shrek', 'Pat Solitano', 'Karl Childers', 'Peter Parker', 'Bruno Antony', 'Seth', 'Caden Cotard', 'Travis Bickle', 'Stanley Ipkiss', 'Raylan Givens', 'Lyn Cassady', 'Michael Scott', 'Robert Angier', 'Rachel Lang', 'Dr. Frank-N-Furter', 'Jack Torrance', 'Tom Ripley', 'D_Artagnan', 'Stephen Hawking', 'Thor', 'James Bond', 'Mark Renton', 'Tugg Speedman', 'David Aames', 'Rorschach', 'Jordan Belfort', 'Logan', 'Judy Hoops', 'Blair Waldorf']
role_names = role_names

password_hash = "652c7dc687d98c9889304ed2e408c74b611e86a40caa51c4b43f1dd5913c5cd0"

def check_password( password ):
    return hashlib.sha256(password.encode('utf-8')).hexdigest() == password_hash

NAME_DICT = {}

for role_name in role_names:
  NAME_DICT[role_name] = role_name

ai_roles_obj = {}

images = {}

for role_name in NAME_DICT:
    role_en = NAME_DICT[role_name]
    img_path = f'images/{role_en}.jpg'
    if os.path.exists(img_path):
        print(f"{img_path} loaded.")
        images[role_en] = Image.open(img_path)
    else:
        print(f"{img_path} not found.")
        images[role_en] = None


for role_name in tqdm(role_names):
    hf_name = 'silk-road/ChatHaruhi-from-RoleLLM/' + role_name
    ai_roles_obj[role_name] = ChatHaruhi( role_from_hf = hf_name, \
                      llm = 'openai', \
                      embedding = 'bge_en')

async def update_image(password, input_image, image_role_name):
    if not check_password(password):
        return "wrong password"

    image_role_name = image_role_name.strip()

    if image_role_name not in role_names:
        return image_role_name + " not in list! check your name!"

    if input_image is not None:
        images[image_role_name] = input_image
        return "update " + image_role_name + " 's photo!"

    

async def get_response(user_role, user_text, ai_role, chatbot):
    role_en = NAME_DICT[ai_role]
    ai_roles_obj[role_en].dialogue_history = copy.deepcopy(chatbot)
    response = ai_roles_obj[role_en].chat(role=user_role, text=user_text)
    user_msg = user_role + ':「' + user_text + '」'
    chatbot.append((user_msg, response))
    print((user_msg, response))
    return chatbot


async def respond(user_role, user_text, ai_role, chatbot):
    return await get_response(user_role, user_text, ai_role, chatbot), None


def clear(user_role, user_text, chatbot):
    return None, None, []


def get_image(ai_role):
    role_en = NAME_DICT[ai_role]
    img_ans = images[role_en]

    return img_ans, "New Friend", None, []


with gr.Blocks() as demo:
    gr.Markdown(
        """
        # Chat凉宫春日 ChatHaruhi (English Version)
        ## Reviving Anime Character in Reality via Large Language Model

        95 English Roles Adapated from [RoleLLM](https://github.com/InteractiveNLP-Team/RoleLLM-public)

        ChatHaruhi2.0的demo implemented by Cheng Li

        More Information See in [https://github.com/LC1332/Chat-Haruhi-Suzumiya](https://github.com/LC1332/Chat-Haruhi-Suzumiya/blob/main/README_EN.md)

        If you find it interesting, please be kind enough to give us a star. 如果觉得有趣请拜托为我们点上star.

        user_role is the character that the person is roleplaying as. Please try to set it to a character relevant to the storyline as much as possible, and do not use the same name as the protagonist.

        using BGE Small English as embedding model other language should unsupported(or wrong memory retrieve)
        """
    )
    with gr.Tab("ChatHaruhi-English"):
        with gr.Row():
            chatbot = gr.Chatbot()
            role_image = gr.Image(height=400, value=None)
        with gr.Row():
            user_role = gr.Textbox(label="user_role", scale = 1, interactive = True, value = "New Friend")
            user_text = gr.Textbox(label="user_text", scale = 20)
        with gr.Row():
            submit = gr.Button("Submit")
            clean = gr.ClearButton(value="Clear")
        ai_role = gr.Radio(role_names, label="characters", value='Sherlock Holmes')

    with gr.Tab("Change Image"):
        with gr.Row():
            idle_textbox = gr.Textbox(value = "Since there are 95 characters, it is a bit difficult to find images for all of them at once. If character you like currently does not have an image, please contact us to upload one.")
        with gr.Row():
            password = gr.Textbox(label="write original password string")

        with gr.Row():
            input_image = gr.Image(interactive = True)

        with gr.Row():
            image_role_name = gr.Textbox(label="Role Name ")

        with gr.Row():
            image_submit = gr.Button("Submit")

    image_submit.click(fn = update_image, inputs = [password, input_image, image_role_name], outputs = [idle_textbox])

    
    ai_role.change(get_image, ai_role, [role_image, user_role, user_text, chatbot])
    user_text.submit(fn=respond, inputs=[user_role, user_text, ai_role, chatbot], outputs=[chatbot, user_text])
    submit.click(fn=respond, inputs=[user_role, user_text, ai_role, chatbot], outputs=[chatbot, user_text])
    clean.click(clear, [user_role, user_text, chatbot], [user_role, user_text, chatbot])
demo.launch()