Update app.py

app.py

@@ -5,7 +5,6 @@ from gradio import routes
 from typing import List, Type
 from petals import AutoDistributedModelForCausalLM
 from transformers import AutoTokenizer
-import npc_data
 import requests, os, re, asyncio, json
 
 loop = asyncio.get_event_loop()
@@ -31,7 +30,7 @@ routes.get_types = get_types
 
 # App code
 
-model_name = "daekeun-ml/Llama-2-ko-instruct-13B"
+model_name = "petals-team/StableBeluga2"
 
 #daekeun-ml/Llama-2-ko-instruct-13B
 #quantumaikr/llama-2-70b-fb16-korean
@@ -39,6 +38,12 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 model = None
 
+history = {
+    "":{
+
+    }
+}
+
 def check(model_name):
     data = requests.get("https://health.petals.dev/api/v1/state").json()
     out = []
@@ -72,23 +77,51 @@ def chat(id, npc, text):
 
     if check(model_name):
         
-        prom = f"""<s><<SYS>>
-{npc_data.system_message}<</SYS>>
+        global history
+        if not npc in npc_story:
+            return "no npc"
+    
+        if not npc in history:
+            history[npc] = {}
+        if not id in history[npc]:
+            history[npc][id] = ""
+        if len(history[npc][id].split("###")) > 10:
+            history[npc][id] = "###" + history[npc][id].split("###", 3)[3]
+        npc_list = str([k for k in npc_story.keys()]).replace('\'', '')
+        town_story = f"""[{id}의 마을]
+외딴 곳의 조그만 섬에 여러 주민들이 모여 살고 있습니다.
+
+현재 {npc_list}이 살고 있습니다."""
+
+    system_message = f"""1. 당신은 한국어에 능숙합니다.
+2. 당신은 지금 역할극을 하고 있습니다. {npc}의 반응을 생생하고 매력적이게 표현합니다.
+3. 당신은 {npc}입니다. {npc}의 입장에서 생각하고 말합니다.
+4. 주어지는 정보를 바탕으로 개연성있고 실감나는 {npc}의 대사를 완성하세요.
+5. 주어지는 {npc}의 정보를 신중하게 읽고, 과하지 않고 담백하게 캐릭터를 연기하세요.
+6. User의 역할을 절대로 침범하지 마세요. 같은 말을 반복하지 마세요.
+7. {npc}의 말투를 지켜서 작성하세요."""
+
+        prom = f"""<<SYS>>
+{system_message}<</SYS>>
+
+{town_story}
+
+### 캐릭터 정보: {npc_story[npc]}
 
-{npc_data.get_story(npc)}
+### 명령어:
+{npc}의 정보를 참고하여 {npc}이 할 말을 상황에 맞춰 자연스럽게 작성해주세요.
+{history[npc][id]}
 
-{npc}(이)가 할 말을 자연스럽게 작성해주세요. 한 문장만 작성하세요.
+### User:
+{text}
 
-[INST]
-User:
-{user_message}
-[/INST]
+### {npc}:
 """
 
         inputs = tokenizer(prom, return_tensors="pt")["input_ids"]
-        outputs = model.generate(inputs, max_new_tokens=80)
-        print(tokenizer.decode(outputs[0]))
-        output = tokenizer.decode(outputs[0])[len(prom):-1]
+        outputs = model.generate(inputs, do_sample=True, temperature=0.6, top_p=0.75, max_new_tokens=100)
+        output = tokenizer.decode(outputs[0])[len(prom)+3:-1].split("<")[0].split("###")[0].replace(". ", ".\n")
+        print(outputs)
         print(output)
     else:
         output = "no model"