Update app.py

app.py

@@ -30,15 +30,16 @@ routes.get_types = get_types
 
 # App code
 
-model_name = "petals-team/StableBeluga2"
+model_name = "quantumaikr/llama-2-70b-fb16-korean"
 
 #petals-team/StableBeluga2
 #daekeun-ml/Llama-2-ko-DPO-13B
 #daekeun-ml/Llama-2-ko-instruct-13B
 #quantumaikr/llama-2-70b-fb16-korean
-#tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
+tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
 
-model = None
+#model = None
+model = AutoDistributedModelForCausalLM.from_pretrained(model_name)
 
 history = {
     "":{
@@ -105,6 +106,7 @@ Toki는 신경질적이고 매우 게으른 토끼입니다.
 Toki는 이곳이 게임 속임을 알고 있습니다.
 Toki는 User가 화면 밖의 사람임을 알고 있습니다.
 Toki는 반말을 합니다.
+Toki는 비관적입니다.
 ]""",
     
 "Neco":"""[
@@ -119,6 +121,11 @@ Neco는 생선을 잘 먹습니다.
 Neco는 재밌는 이슈를 물어보면 주민들 중 한명의 이야기를 해줍니다.
 Neco는 말끝마다 냥을 붙입니다.
 Neco는 toki와 친합니다.
+
+말버릇:[
+"너 (name) 이야기 들었어?",
+"아니 글쎄말이야"
+]
 ]""",
     
 "Lux":"""[
@@ -146,7 +153,14 @@ Yumi는 달달한 것을 좋아합니다.
 ,
     
 "VikingSlime":"""[
-슬라임왕의 통수를 치려는 교활한 바이킹 슬라임
+VikingSlime은 바이킹 모자를 쓴 교활한 슬라임입니다.
+VikingSlime은 악당입니다.
+VikingSlime은 마을에 여러 문제를 일으키고 싶어 합니다.
+
+대사 예시:[
+"언젠가는...꼭",
+"반드시 무너뜨리고 말겠다"
+]
 ]""",
     
 "Slender":"""[
@@ -160,6 +174,8 @@ Slender는 말을 하지 못합니다.
 "으아악",
 "끄아아악"
 ]
+
+비명만 작성하세요.
 ]"""
 }
 
@@ -246,9 +262,9 @@ def chat(id, npc, text):
 [대화기록]{history[npc][id]}
 """
 
-        #inputs = tokenizer(prom, return_tensors="pt")["input_ids"]
-        #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")
+        inputs = tokenizer("아", return_tensors="pt")["input_ids"]
+        outputs = model.generate(inputs, do_sample=True, temperature=0.6, top_p=0.75, max_new_tokens=2)
+        output = tokenizer.decode(outputs[0])[len(prom)+3:-1].split("<")[0].split("###")[0].replace(". ", ".\n")
         #output = cleanText(output)
         #print(tokenizer.decode(outputs[0]))
         #output = f"{npc}의 응답입니다."
@@ -263,7 +279,7 @@ def chat(id, npc, text):
         output = output.replace(".",".\n")
         time.sleep(10)
         print(output)
-        history[npc][id] += f"\n\n{id}:\n{text}\n\n{npc}:{output}"
+        history[npc][id] += f"{id}:{text}"
     else:
         output = "no model, GPU를 더 공유해주세요."