maybe complete

app.py

@@ -1,144 +1,114 @@
 import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import warnings
+
+
+class Chatbot():
+    def __init__(self):
+        self.tokenizer = AutoTokenizer.from_pretrained('kakaobrain/kogpt', revision='KoGPT6B-ryan1.5b')
+        special_tokens_dict = {'additional_special_tokens': ['<sep>', '<eos>', '<sos>', '#@이름#', '#@계정#', '#@신원#', '#@전번#', '#@금융#', '#@번호#', '#@주소#', '#@소속#', '#@기타#']}
+        num_added_toks = self.tokenizer.add_special_tokens(special_tokens_dict)
+
+        self.model = AutoModelForCausalLM.from_pretrained("/workspace/test_trainer/checkpoint-10000")
+        self.model.resize_token_embeddings(len(self.tokenizer))
+        self.model = self.model.cuda()
+
+        self.info = None
+        self.talk = []
+
+    def initialize(self, topic, bot_addr, bot_age, bot_sex, my_addr, my_age, my_sex):
+        def encode(age):
+            if age < 20:
+                age = "20대 미만"
+            elif age >= 70:
+                age = "70대 이상"
+            else:
+                age = str(age // 10 * 10) + "대"
+            return age
+        bot_age = encode(bot_age)
+        my_age = encode(my_age)
+        self.info = f"일상 대화 {topic}<sep>P01:{my_addr} {my_age} {my_sex}<sep>P02:{bot_addr} {bot_age} {bot_sex}<sep>"
+        return self.info_check()
+
+    def info_check(self):
+        return self.info.replace('<sep>', '\n').replace('P01', '당신').replace('P02', '챗봇')
+
+    def reset_talk(self):
+        self.talk = []
+
+    def test(self, myinp):
+        state = None
+        inp = "P01<sos>" + myinp + "<eos>"
+        self.talk.append(inp)
+        self.talk.append("P02<sos>")
 
-def update(name):
-    return f"Welcome to Gradio, {name}!"
-
-demo = gr.Blocks()
-
-with demo:
-    gr.Markdown(f"각 질문에 대답 후 Enter 해주세요.\n\n")
-    with gr.Row():
-        topic = gr.Textbox(label="Topic", placeholder="대화 주제를 정해주세요 (e.g. 여가 생활, 일과 직업, 개인 및 관계, etc...)")
-    with gr.Row():
-        with gr.Column():
-            addr = gr.Textbox(label="지역", placeholder="e.g. 여가 생활, 일과 직업, 개인 및 관계, etc...")
-            age = gr.Textbox(label="나이", placeholder="e.g. 20대 미만, 40대, 70대 이상, etc...")
-            sex = gr.Textbox(label="성별", placeholder="e.g. 남성, 여성, etc...")
-        with gr.Column():
-            addr = gr.Textbox(label="지역", placeholder="e.g. 여가 생활, 일과 직업, 개인 및 관계, etc...")
-            age = gr.Textbox(label="나이", placeholder="e.g. 20대 미만, 40대, 70대 이상, etc...")
-            sex = gr.Textbox(label="성별", placeholder="e.g. 남성, 여성, etc...")
-        out = gr.Textbox()
-    btn = gr.Button("Run")
-    # btn.click(fn=update, inputs=inp, outputs=out)
-
-demo.launch()
+        while True:
+            now_inp = self.info + "".join(self.talk)
+            inputs = self.tokenizer(now_inp, max_length=1024, truncation='longest_first', return_tensors='pt')
+            seq_len = inputs.input_ids.size(1)
+            if seq_len > 512 * 0.8:
+                state = f"<주의> 현재 대화 길이가 곧 최대 길이에 도달합니다. ({seq_len} / 512)"
+
+            if seq_len >= 512:
+                state = "<주의> 대화 길이가 너무 길어졌기 때문에, 이후 대화는 맨 앞의 발화를 조금씩 지우면서 진행됩니다."
+                talk = talk[1:]
+            else:
+                break
+            
+        out = self.model.generate(
+            inputs=inputs.input_ids.cuda(), 
+            attention_mask=inputs.attention_mask.cuda(),
+            max_length=512, 
+            do_sample=True,
+            pad_token_id=self.tokenizer.pad_token_id,
+            eos_token_id=self.tokenizer.encode('<eos>')[0]
+            )
+        out = self.tokenizer.batch_decode(out)
+        real_out = out[0][len(now_inp):-5]
+        self.talk[-1] += out[0][len(now_inp):]
+        return [(self.talk[i][8:-5], self.talk[i+1][8:-5]) for i in range(0, len(self.talk)-1, 2)]
 
 
-def main(model_name):
+if __name__ == "__main__":
     warnings.filterwarnings("ignore")
 
-    tokenizer = AutoTokenizer.from_pretrained('kakaobrain/kogpt', revision='KoGPT6B-ryan1.5b')
-    special_tokens_dict = {'additional_special_tokens': ['<sep>', '<eos>', '<sos>', '#@이름#', '#@계정#', '#@신원#', '#@전번#', '#@금융#', '#@번호#', '#@주소#', '#@소속#', '#@기타#']}
-    num_added_toks = tokenizer.add_special_tokens(special_tokens_dict)
-
-    model = AutoModelForCausalLM.from_pretrained(model_name)
-    model.resize_token_embeddings(len(tokenizer))
-    model = model.cuda()
-
-    info = ""
-    
-    while True:
-        if info == "":
-            print(
-                f"지금부터 대화 정보를 입력 받겠습니다.\n"
-                f"각 질문에 대답 후 Enter 해주세요.\n"
-                f"아무 입력 없이 Enter 할 경우, 미리 지정된 값 중 랜덤으로 정하게 됩니다.\n"
-            )
+    chatbot = Chatbot()
+    demo = gr.Blocks()
+
+    with demo:
+        gr.Markdown("# <center>MINDs Lab Brain's Fast Neural Chit-Chatbot</center>")
+        with gr.Row():
+            with gr.Column():
+                topic = gr.Radio(label="Topic", choices=['여가 생활', '시사/교육', '미용과 건강', '식음료', '상거래(쇼핑)', '일과 직업', '주거와 생활', '개인 및 관계', '행사'])
+                with gr.Column():
+                    gr.Markdown(f"Bot's persona")
+                    bot_addr = gr.Dropdown(label="지역", choices=['서울특별시', '경기도', '부산광역시', '대전광역시', '광주광역시', '울산광역시', '경상남도', '인천광역시', '충청북도', '제주도', '강원도', '충청남도', '전라북도', '대구광역시', '전라남도', '경상북도', '세종특별자치시', '기타'])
+                    bot_age = gr.Slider(label="나이", minimum=10, maximum=80, value=45, step=1)
+                    bot_sex = gr.Radio(label="성별", choices=["남성", "여성"])
+                with gr.Column():
+                    gr.Markdown(f"Your persona")
+                    my_addr = gr.Dropdown(label="지역", choices=['서울특별시', '경기도', '부산광역시', '대전광역시', '광주광역시', '울산광역시', '경상남도', '인천광역시', '충청북도', '제주도', '강원도', '충청남도', '전라북도', '대구광역시', '전라남도', '경상북도', '세종특별자치시', '기타'])
+                    my_age = gr.Slider(label="나이", minimum=10, maximum=80, value=45, step=1)
+                    my_sex = gr.Radio(label="성별", choices=["남성", "여성"])
+                with gr.Row():
+                    btn = gr.Button(label="적용")
+                    state = gr.Textbox(label="상태")
+                    btn.click(
+                        fn=chatbot.initialize, 
+                        inputs=[topic, bot_addr, bot_age, bot_sex, my_addr, my_age, my_sex], 
+                        outputs=state
+                    )
+                    
+            with gr.Column():
+                screen = gr.Chatbot(label="익명의 상대")
+                with gr.Row():
+                    speak = gr.Textbox(label="입력창")
+                    btn = gr.Button(label="Talk")
+                    btn.click(
+                        fn=chatbot.test, 
+                        inputs=speak, 
+                        outputs=screen
+                    )
+    demo.launch(share=True)
 
-            time.sleep(1)
 
-            yon = "no"
-        else:
-            yon = input(
-                f"이전 대화 정보를 그대로 유지할까요? (yes : 유지, no : 새로 작성) :"
-            )
-        
-        if yon == "no":
-            info = "일상 대화 "
-
-            topic = input("대화 주제를 정해주세요 (e.g. 여가 생활, 일과 직업, 개인 및 관계, etc...) :")
-            if topic == "":
-                topic = random.choice(['여가 생활', '시사/교육', '미용과 건강', '식음료', '상거래(쇼핑)', '일과 직업', '주거와 생활', '개인 및 관계', '행사'])
-                print(topic)
-            info += topic + "<sep>"
-
-            def ask_info(who, ment):
-                print(ment)
-                text = who + ":"
-                addr = input("어디 사세요? (e.g. 서울특별시, 제주도, etc...) :").strip()
-                if addr == "":
-                    addr = random.choice(['서울특별시', '경기도', '부산광역시', '대전광역시', '광주광역시', '울산광역시', '경상남도', '인천광역시', '충청북도', '제주도', '강원도', '충청남도', '전라북도', '대구광역시', '전라남도', '경상북도', '세종특별자치시', '기타'])
-                    print(addr)
-                text += addr + " "
-
-                age = input("나이가? (e.g. 20대, 70대 이상, etc...) :").strip()
-                if age == "":
-                    age = random.choice(['20대', '30대', '50대', '20대 미만', '60대', '40대', '70대 이상'])
-                    print(age)
-                text += age + " "
-
-                sex = input("성별이? (e.g. 남성, 여성, etc... (?)) :").strip()
-                if sex == "":
-                    sex = random.choice(['남성', '여성'])
-                    print(sex)
-                text += sex + "<sep>"
-                return text
-
-            info += ask_info(who="P01", ment=f"\n당신에 대해 알려주세요.\n")
-            info += ask_info(who="P02", ment=f"\n챗봇에 대해 알려주세요.\n")
-
-        pp = info.replace('<sep>', '\n')
-        print(
-            f"\n----------------\n"
-            f"<입력 정보 확인> (P01 : 당신, P02 : 챗봇)\n"
-            f"{pp}"
-            f"----------------\n"
-            f"대화를 종료하고 싶으면 언제든지 'end' 라고 말해주세요~\n"
-        )
-        talk = []
-        switch = True
-        switch2 = True
-        while True:
-            inp = "P01<sos>"
-            myinp = input("당신 : ")
-            if myinp == "end":
-                print("대화 종료!")
-                break
-            inp += myinp + "<eos>"
-            talk.append(inp)
-            talk.append("P02<sos>")
-
-            while True:
-                now_inp = info + "".join(talk)
-                inpu = tokenizer(now_inp, max_length=1024, truncation='longest_first', return_tensors='pt')
-                seq_len = inpu.input_ids.size(1)
-                if seq_len > 512 * 0.8 and switch:
-                    print(
-                        f"<주의> 현재 대화 길이가 곧 최대 길이에 도달합니다. ({seq_len} / 512)"
-                    )
-                    switch = False
-
-                if seq_len >= 512 and switch2:
-                    print("<주의> 대화 길이가 너무 길어졌기 때문에, 이후 대화는 맨 앞의 발화를 조금씩 지우면서 진행됩니다.")
-                    talk = talk[1:]
-                    switch2 = False
-                else:
-                    break
-            
-            out = model.generate(
-                inputs=inpu.input_ids.cuda(), 
-                attention_mask=inpu.attention_mask.cuda(),
-                max_length=512, 
-                do_sample=True,
-                pad_token_id=tokenizer.pad_token_id,
-                eos_token_id=tokenizer.encode('<eos>')[0]
-                )
-            output = tokenizer.batch_decode(out)
-            print("챗봇 : " + output[0][len(now_inp):-5])
-            talk[-1] += output[0][len(now_inp):]
-
-        again = input(f"다른 대화를 시작할까요? (yes : 새로운 시작, no : 종료) :")
-        if again == "no":
-            break
-        

flagged/log.csv

@@ -0,0 +1,4 @@
+'self','output','flag','username','timestamp'
+'뭐하고 계세요?','[(''안녕하세요'', ''넵''), (''뭐하고 계세요?'', ''저 게임하면서 있어용'')]','','','2022-06-29 07:59:03.609856'
+'뭐하고 계세요?','[(''안녕하세요'', ''넵''), (''뭐하고 계세요?'', ''저 게임하면서 있어용'')]','','','2022-06-29 07:59:07.265460'
+'안녕하세요?','[[''안녕하세요?'', ''아니'']]','','','2022-06-29 08:15:33.284872'