generation parameter 변경

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
 from streamlit_chat import message
 
+
 @st.cache(allow_output_mutation=True)
 def get_pipe():
     from transformers import AutoTokenizer, AutoModelForCausalLM
@@ -10,35 +11,38 @@ def get_pipe():
     return model, tokenizer
 
 def get_response(tokenizer, model, history, max_context: int = 7, bot_id: str = '1'):
+    # print("history:", history)
     context = []
     for i, text in enumerate(history):
-        context.append(f"{i % 2} : {text}</s>")
+        context.append(f"{i % 2}: {text}</s>")
     
     if len(context) > max_context:
         context = context[-max_context:]
-    context = "".join(context) + f"{bot_id} : "
+    context = "".join(context) + f"{bot_id}: "
     inputs = tokenizer(context, return_tensors="pt")
     
     generation_args = dict(
-        max_new_tokens=64,
+        max_new_tokens=128,
         min_length=inputs["input_ids"].shape[1] + 5,
+        # no_repeat_ngram_size=4,
         eos_token_id=2,
         do_sample=True,
-        top_p=0.6,
-        temperature=0.8,
-        repetition_penalty=1.5,
+        top_p=0.95,
+        temperature=1.35,
+        # repetition_penalty=1.0,
         early_stopping=True
     )
 
     outputs = model.generate(**inputs, **generation_args)
-    response = tokenizer.decode(outputs[0])
-    print(context)
-    print(response)
-    response = response[len(context):].replace("</s>", "")
-
+    response = tokenizer.decode(outputs[0], skip_special_tokens=False)
+    print("Context:", tokenizer.decode(inputs["input_ids"][0]))
+    print("Response:", response)
+    response = response[len(context):].replace("</s>", "").replace("\n", "")
+    response = response.split("<s>")[0]
+    # print("Response:", response)
     return response
 
-st.title("한국어 대화 모델 demo")
+st.title("ajoublue-gpt2-medium 한국어 대화 모델 demo")
 
 with st.spinner("loading model..."):
     model, tokenizer = get_pipe()