Add application file

app.py

@@ -0,0 +1,436 @@
+from sortinghat import SortingHat
+import random
+import torch
+import gradio as gr
+from PIL import Image
+
+from peft import AutoPeftModelForCausalLM
+
+from transformers import AutoTokenizer
+from time import sleep
+
+class HPChatBot:
+    def __init__(self,
+                 model_path: str = 'rnltls/harrypotter_lexicon_finetune',
+                 device_map: str = 'auto',
+                 load_in_4_bit: bool = True,
+                 **quant_kwargs) -> None:
+        self.model = None
+        self.tokenizer = None
+        self.image_processor = None
+        self.conv = None
+        self.conv_img = []
+        self.img_tensor = []
+        self.roles = None
+        self.stop_key = None
+        self.is_chat = False
+        self.is_waldo = False
+        self.load_models(model_path,
+                         device_map=device_map,
+                         load_in_4_bit=load_in_4_bit,
+                         **quant_kwargs)
+
+    def load_models(self, model_path: str,
+                    device_map: str,
+                    load_in_4_bit: bool,
+                    **quant_kwargs) -> None:
+        model = AutoPeftModelForCausalLM.from_pretrained(
+            model_path, # YOUR MODEL YOU USED FOR TRAINING
+            load_in_4bit = load_in_4_bit,
+        )
+        tokenizer = AutoTokenizer.from_pretrained(model_path)
+
+        self.model = AutoPeftModelForCausalLM.from_pretrained(
+            "rnltls/harrypotter_lexicon_finetune", # YOUR MODEL YOU USED FOR TRAINING
+            load_in_4bit = load_in_4_bit,
+        )
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path,
+                                                       use_fast=False)
+    
+    def generate_answer(self, prompt):
+        output = self.model.generate(**prompt, max_new_tokens = 256)
+        generated_text = self.tokenizer.decode(output[0], skip_special_tokens=True)
+        return generated_text
+
+HogwartGemma = "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/HogwartGemma.jpg?raw=true"
+SpellProf = ["https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/Spell_stand.jpg?raw=true",
+           "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/Spell_think.jpg?raw=true",
+           "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/Spell_show.jpg?raw=true"]
+PotionProf = ["https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/Potion_stand.jpg?raw=true",
+              "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/Potion_think.jpg?raw=true",
+              "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/Potion_show.jpg?raw=true"]
+OtherProf = ["https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/Other_stand.jpg?raw=true",
+             "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/Other_think.jpg?raw=true",
+             "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/Other_show.jpg?raw=true"]
+HouseHat = ["https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/sorting_hat.jpg?raw=true",
+          "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/sorting_hat_think.jpg?raw=true",
+          "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/sorting_hat_gryffindor.png?raw=true",
+            "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/sorting_hat_hufflepuff.png?raw=true",
+            "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/sorting_hat_ravenclaw.png?raw=true",
+          "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/sorting_hat_slytherin.png?raw=true"]
+Icons = ["https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/wand_icon.png?raw=true",
+         "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/cauldron.png?raw=true",
+         "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/golden-snitch.png?raw=true",
+         "https://github.com/kangshwan/hogwarts_chatbot/blob/main/IMG/hat_icon.png?raw=true"]
+enable_btn = gr.Button(interactive=True)
+disable_btn = gr.Button(interactive=False)
+spell_chat, potion_chat, other_chat, house_chat = False, False, False, False
+question_counter = 0
+quit_counter = 0
+
+question_list = None
+quote_list = None
+answer_list = []
+question_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
+
+### Instruction:
+{}
+
+### Input:
+{}
+
+### Response:
+{}"""
+
+def show_image(path):
+    # Convert To PIL Image
+    image = Image.open(path)
+    return image
+
+def spell_chatting(chat_history, request: gr.Request):
+    global spell_chat, potion_chat, other_chat, house_chat, quit_counter, answer_list
+    answer_list = []
+    quit_counter=0
+    question_counter = 0
+    chat_history.clear()
+    spell_chat, potion_chat, other_chat, house_chat = True, False, False, False
+    bot_message = "Welcome to Magic Spell Class!\nTell me what you want to achieve, and I’ll suggest the perfect spell for it!\nIf you ask in the format: 'What spell can I use when I ~?', I can give you even better suggestions!"
+    txt_box = gr.Textbox(value="What spell can I use when I ", interactive=True)
+    prof_IMG = SpellProf[0]
+    chat_history.append([None, bot_message])
+    return chat_history, prof_IMG, txt_box, enable_btn, enable_btn, enable_btn
+
+def potion_chatting(chat_history, request: gr.Request):
+    global spell_chat, potion_chat, other_chat, house_chat, quit_counter, answer_list
+    answer_list = []
+    quit_counter=0
+    question_counter = 0
+    chat_history.clear()
+    spell_chat, potion_chat, other_chat, house_chat = False, True, False, False
+    bot_message = "Welcome to Potion Class!\nTell me what you want to achieve, and I’ll suggest the perfect potion for it!\nIf you ask in the format: 'What potion can I make when I ~?', I can give you even better suggestions!"
+    txt_box = gr.Textbox(value="What potion can I make when I ", interactive=True)
+    prof_IMG = PotionProf[0]
+    chat_history.append([None,bot_message])
+    return chat_history, prof_IMG, txt_box, enable_btn, enable_btn, enable_btn
+
+def other_chatting(chat_history, request: gr.Request):
+    global spell_chat, potion_chat, other_chat, house_chat, quit_counter, answer_list
+    answer_list = []
+    quit_counter=0
+    question_counter = 0
+    chat_history.clear()
+    spell_chat, potion_chat, other_chat, house_chat = False, False, True, False
+    bot_message = "Welcome to Library!\nWelcome to the library! Feel free to ask me anything if you're curious!"
+    txt_box = gr.Textbox(placeholder="Ask anything you're curious about in the Wizarding World!", value="", interactive=True)
+    prof_IMG = OtherProf[0]
+    chat_history.append([None,bot_message])
+    return chat_history, prof_IMG, txt_box, enable_btn, enable_btn, enable_btn
+
+def house_chatting(chat_history, request: gr.Request):
+    global spell_chat, potion_chat, other_chat, house_chat, quit_counter, question_counter, question_list, quote_list, answer_list
+    answer_list = []
+    question_list = sortinghat.choose_questions()
+    quote_list = sortinghat.choose_quotes()
+    quit_counter=0
+    question_counter = 0
+    chat_history.clear()
+    spell_chat, potion_chat, other_chat, house_chat = False, False, False, True
+    bot_message = "Hmm, let's see… where shall I place you?\n" + question_list[question_counter]
+    txt_box = gr.Textbox(placeholder="Answer the question", value="", interactive=True)
+    prof_IMG = HouseHat[0]
+    chat_history.append([None,bot_message])
+    return chat_history, prof_IMG, txt_box, enable_btn, enable_btn, enable_btn
+
+def ask_question(chat_history, text_data, request: gr.Request):
+    global spell_chat, potion_chat, other_chat, house_chat, quit_counter, answer_list
+    if spell_chat:
+        print(text_data)
+        prof_IMG = SpellProf[1]
+        chat_history.append([text_data, None])
+        text_data = gr.Textbox(value="What spell can I use when I ", interactive=True)
+        return chat_history, prof_IMG, text_data
+
+    if potion_chat:
+        print(text_data)
+        prof_IMG = PotionProf[1]
+        chat_history.append([text_data, None])
+        text_data = gr.Textbox(value="What potion can I make when I ", interactive=True)
+        return chat_history, prof_IMG, text_data
+
+    if other_chat:
+        print(text_data)
+        prof_IMG = OtherProf[1]
+        chat_history.append([text_data, None])
+        text_data = gr.Textbox(placeholder="Ask anything you're curious about in the Wizarding World!", value="", interactive=True)
+        return chat_history, prof_IMG, text_data
+
+    if house_chat:
+        print(text_data)
+        prof_IMG = HouseHat[1]
+        chat_history.append([text_data, None])
+        answer_list.append(text_data)
+        text_data = gr.Textbox(placeholder="Answer the question", value = "", interactive=True)
+        return chat_history, prof_IMG, text_data
+
+def clean_chatting(chat_history, request: gr.Request):
+    global spell_chat, potion_chat, other_chat, house_chat, quit_counter, question_counter, question_list, quote_list, answer_list
+    quit_counter = 0
+    question_counter = 0
+    question_list=sortinghat.choose_questions()
+    quote_list=sortinghat.choose_quotes()
+    answer_list=[]
+    chat_history.clear()
+    if spell_chat:
+        bot_message = "Welcome to Magic Spell Class!\nTell me what you want to achieve, and I’ll suggest the perfect spell for it!\nIf you ask in the format: 'What spell can I use when I ~?', I can give you even better suggestions!"
+        prof_IMG = SpellProf[0]
+    if potion_chat:
+        bot_message = "Welcome to Potion Class!\nTell me what you want to achieve, and I’ll suggest the perfect potion for it!\nIf you ask in the format: 'What potion can I make when I ~?', I can give you even better suggestions!"
+        prof_IMG = PotionProf[0]
+    if other_chat:
+        bot_message = "Welcome to Library!\nWelcome to the library! Feel free to ask me anything if you're curious!"
+        prof_IMG = OtherProf[0]
+    if house_chat:
+        bot_message = "Hmm, let's see… where shall I place you?\n" + question_list[question_counter]
+        prof_IMG = HouseHat[0]
+    chat_history.append([None, bot_message])
+    return chat_history, prof_IMG
+
+def run_model(chat_history, request: gr.Request):
+    global spell_chat, potion_chat, other_chat, house_chat, question_counter, question_list, quote_list, answer_list
+
+    if question_counter < 4 and house_chat:
+        bot_message = quote_list[question_counter] + '\n' + question_list[question_counter+1]
+        question_counter += 1
+        chat_history[-1][1] = bot_message
+        return chat_history
+
+    elif house_chat:
+        print("Choosing the house")
+        print(question_list)
+        print(answer_list)
+        user_responses = "\n".join([f"{i+1}. {q}\nAnswer: {a}" for i, (q,a) in enumerate(zip(question_list, answer_list))])
+
+        prompt = f"""
+            You are the Sorting Hat at Hogwarts. Based on the student's answers to the following questions, assign them to the most suitable house (Gryffindor, Hufflepuff, Ravenclaw, or Slytherin) and provide three reasons for your decision.
+            Gryffindor values courage, nerve, and chivalry.
+            Hufflepuff values hard work, patience, justice, and loyalty.
+            Ravenclaw values intelligence, learning, wisdom, and wit.
+            Slytherin values ambition, cunning, leadership, and resourcefulness.
+
+            ### Instruction:
+            Please format the response like this:
+            'Your house is [house]!
+            Here's why:
+            1. [reason 1]
+            2. [reason 2]
+            3. [reason 3]
+            Welcome to your new house at Hogwarts!'
+
+            ### Input:
+            {user_responses}
+
+            ### Response:
+        """
+        inputs = HP_chatbot.tokenizer([prompt], return_tensors="pt").to("cuda")
+        output = HP_chatbot.generate_answer(inputs)
+        output = output.split("Response:")[-1]
+        print(output)
+        chat_history[-1][1] = output
+        return chat_history
+
+    alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
+
+    ### Instruction:
+    {}
+
+    ### Input:
+    {}
+
+    ### Response:
+    {}"""
+
+    inputs = HP_chatbot.tokenizer(
+    [
+        alpaca_prompt.format(
+            chat_history[-1][0], # instruction
+            "", # input
+            "", # output - leave this blank for generation!
+        )
+    ], return_tensors = "pt").to("cuda")
+    output = HP_chatbot.generate_answer(inputs)
+
+    output = output.split("Response:")[-1]
+    chat_history[-1][1] = output
+    return chat_history
+
+def end_chatting(chat_history, textbox, request: gr.Request):
+    global spell_chat, potion_chat, other_chat, house_chat, quit_counter, question_counter
+    if not house_chat or quit_counter > 0:
+        quit_counter = 0
+        question_counter = 0
+        chat_history.clear()
+        chat_history.append([None, "Welcome to Hogwarts!"])
+        textbox = gr.Textbox(show_label=False, placeholder="Welcome to Hogwart!", value= "", container=False, interactive = False)
+        return chat_history, HogwartGemma, textbox, disable_btn, disable_btn, disable_btn
+    else:
+        quit_counter += 1
+        chat_history.append([None, "Quit, you say? Oh, I see your plight, But the housing question’s still in sight.\
+                             \nThough weary you may be, don't yet retreat,For there's still a task you must complete!"])
+        return chat_history, HouseHat[0], textbox, enable_btn, enable_btn, enable_btn
+
+def result_img(chat_history, textbox, submit_btn,  request: gr.Request):
+    global spell_chat, potion_chat, other_chat, house_chat, quit_counter, question_counter
+    if spell_chat:
+        return SpellProf[2], textbox, submit_btn
+    if potion_chat:
+        return PotionProf[2], textbox, submit_btn
+    if other_chat:
+        return OtherProf[2], textbox, submit_btn
+    if house_chat:
+        if question_counter == 4:
+            quit_counter = 1
+            house = chat_history[-1][1]
+            # print(house)
+            if "gryffindor" in house.lower():
+                insignia =  HouseHat[2]
+            elif "hufflepuff" in house.lower():
+                insignia =  HouseHat[3]
+            elif "ravenclaw" in house.lower():
+                insignia =  HouseHat[4]
+            elif "slytherin" in house.lower():
+                insignia =  HouseHat[5]
+            else:
+                return HouseHat[0], textbox, disable_btn
+            textbox = gr.Textbox(show_label=False, placeholder="Housing Done.", value= "", container=False, interactive = False)
+            return insignia, textbox, disable_btn
+        else:
+            return HouseHat[0], textbox, submit_btn
+
+def build_gradio(concurrency_count=10):
+    textbox = gr.Textbox(show_label=False, placeholder="Welcome to Hogwart!", container=False, interactive = False)
+    with gr.Blocks(
+            theme='gstaff/xkcd',
+        ) as demo:
+        state = gr.State()
+        with gr.Row():
+            with gr.Column(scale=3):
+
+                imagebox = gr.Image(value=HogwartGemma, type="pil", interactive=False, show_label=False, show_download_button=False, show_fullscreen_button=False)
+
+                spell_btn = gr.Button(icon=Icons[0], value="Spell", interactive = True, scale = 2)
+                potion_btn = gr.Button(icon=Icons[1], value="Potion", interactive = True, scale = 2)
+                other_btn = gr.Button(icon=Icons[2], value="Others", interactive = True, scale = 2)
+                house_btn = gr.Button(icon=Icons[3], value="Find your House", interactive = True, scale = 2)
+
+            with gr.Column(scale=8):
+                initial_message = [None,"Welcome to Hogwarts!"]
+                chatbot = gr.Chatbot(
+                    label='Free to ask!',
+                    value= [initial_message],
+                    height=600,
+                )
+                with gr.Row():
+                    with gr.Column(scale=8):
+                        textbox.render()
+                    with gr.Column(scale=1, min_width=50):
+                        submit_btn = gr.Button(value="Send", variant="primary", interactive = False)
+
+                with gr.Row(elem_id="buttons") as button_row:
+                    finish_btn = gr.Button(value="🏁 End Talking", interactive = False)
+                    clear_btn = gr.Button(value="🗑️  Clear", interactive=False)
+
+        box_list = [spell_btn, potion_btn, other_btn, house_btn, textbox, submit_btn, finish_btn, clear_btn]
+
+        spell_btn.click(
+            spell_chatting,
+            inputs=[chatbot],
+            outputs=[chatbot, imagebox, textbox, submit_btn, finish_btn, clear_btn]
+        )
+        potion_btn.click(
+            potion_chatting,
+            inputs=[chatbot],
+            outputs=[chatbot, imagebox, textbox, submit_btn, finish_btn, clear_btn]
+        )
+        other_btn.click(
+            other_chatting,
+            inputs=[chatbot],
+            outputs=[chatbot, imagebox, textbox, submit_btn, finish_btn, clear_btn]
+        )
+        house_btn.click(
+            house_chatting,
+            inputs=[chatbot],
+            outputs=[chatbot, imagebox, textbox, submit_btn, finish_btn, clear_btn]
+        )
+        textbox.submit(
+            ask_question,
+            inputs = [chatbot, textbox],
+            outputs = [chatbot, imagebox, textbox]
+        ).then(
+            run_model,
+            inputs = [chatbot],
+            outputs = [chatbot]
+        ).then(
+            result_img,
+            inputs=[chatbot, textbox, submit_btn],
+            outputs=[imagebox, textbox, submit_btn]
+        )
+
+        submit_btn.click(
+            ask_question,
+            inputs = [chatbot, textbox],
+            outputs = [chatbot, imagebox, textbox]
+        ).then(
+            run_model,
+            inputs = [chatbot],
+            outputs = [chatbot]
+        ).then(
+            result_img,
+            inputs=[chatbot, textbox, submit_btn],
+            outputs=[imagebox, textbox, submit_btn]
+        )
+
+        finish_btn.click(
+            end_chatting,
+            inputs=[chatbot, textbox],
+            outputs=[chatbot, imagebox, textbox, submit_btn, finish_btn, clear_btn]
+        )
+
+        clear_btn.click(
+            clean_chatting,
+            inputs=[chatbot],
+            outputs=[chatbot, imagebox]
+        )
+
+    return demo
+
+import argparse
+if __name__ == "__main__":
+
+    # 컨트롤러에서 사용가능한 모델 가져옴. 나는 모델 하나만 쓸 것이기 때문에, get_model_list()에서 해당 모델이 동작하고 있는 url을 넘겨주면 된다!
+    # models = [args.model_url]
+
+    # HP_chatbot = HPChatBot(load_in_8bit=True,
+    #                    bnb_8bit_compute_dtype=torch.float16,
+    #                    bnb_8bit_use_double_quant=True,
+    #                    bnb_8bit_quant_type='nf8')
+
+    sortinghat = SortingHat()
+
+    # Gradio를 이용해서 데모 만들기
+    demo = build_gradio()
+    demo.queue(
+        api_open=False
+    ).launch(
+        server_port=7861,
+        share=True
+    )

sortinghat.py

@@ -0,0 +1,42 @@
+import random
+class SortingHat:
+    def __init__(self):
+        self.questions = [
+            "How do you approach challenges and obstacles, and what do you think defines bravery?",
+            "When working with others, do you prefer to lead, follow, or collaborate, and why?",
+            "In a difficult situation, would you prioritize being fair or being kind, and how would you balance the two?",
+            "How do you stay motivated when pursuing your goals, and what role does ambition play in your life?",
+            "When faced with rules or expectations, do you tend to follow them, question them, or create your own? Why?",
+            "Which do you value more: being honest or being kind? Why?",
+            "Do you rely more on intuition or analysis when making decisions? Why?",
+            "How do you make sure your ideas and opinions are heard in group discussions?",
+            "What role do you usually take in group projects or collaborations? Why?",
+            "How do you respond to failure, and what does it teach you?",
+            "What principles guide you in making difficult decisions?",
+            "How do you stay grounded while striving for excellence in your work or studies?"
+        ]
+        self.answers = [
+            "So many minds, so many dreams, each one yearning for greatness in their own way.",
+            "Such complexity in this one—more layers than a potion brewed at midnight",
+            "Curious, this mind holds secrets even they haven’t discovered yet.",
+            "I sense a fire within, simmering beneath a calm surface.",
+            "The potential here is boundless, but will they see it themselves?",
+            "Resilient, adaptable, yet a hint of doubt clouds their path.",
+            "A mind as sharp as a blade, cutting through the fog of uncertainty.",
+            "So many roads they could walk, each with its own unique challenge.",
+            "A heart full of dreams and a spirit unbroken—what a combination.",
+            "There’s a quiet strength in this one, hidden but undeniable.",
+            "The choices ahead will shape their destiny more than they know.",
+            "Determined, but with a trace of uncertainty—an intriguing combination.",
+            "They carry a spark of mischief, hidden beneath a composed exterior.",
+            "A mind that seeks adventure, but a heart that yearns for belonging."
+        ]
+
+    def choose_questions(self):
+        self.random_samples = random.sample(self.questions, 5)
+        return self.random_samples
+    
+    def choose_quotes(self):
+        self.random_quotes = random.sample(self.answers, 5)
+        return self.random_quotes
+