simple demo

app.py

@@ -0,0 +1,63 @@
+import streamlit as st
+from streamlit_chat import message
+
+@st.cache_resource()
+def get_pipe():
+    from transformers import AutoTokenizer, AutoModelForCausalLM
+    model_name = "heegyu/ajoublue-gpt2-medium-dialog"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    return model, tokenizer
+
+def get_response(tokenizer, model, history, max_context: int = 7, bot_id: str = '1'):
+    context = []
+    for i, text in enumerate(history):
+        context.append(f"{i % 2} : {text}</s>")
+    
+    if len(context) > max_context:
+        context = context[-max_context:]
+    context = "".join(context) + f"{bot_id} : "
+    inputs = tokenizer(context, return_tensors="pt")
+    
+    generation_args = dict(
+        max_new_tokens=64,
+        min_length=inputs["input_ids"].shape[1] + 5,
+        eos_token_id=2,
+        do_sample=True,
+        top_p=0.6,
+        temperature=0.8,
+        repetition_penalty=1.5,
+        early_stopping=True
+    )
+
+    outputs = model.generate(**inputs, **generation_args)
+    response = tokenizer.decode(outputs[0])
+    print(context)
+    print(response)
+    response = response[len(context):].replace("</s>", "")
+
+    return response
+
+st.title("한국어 대화 모델 demo")
+
+with st.spinner("loading model..."):
+    model, tokenizer = get_pipe()
+
+if 'message_history' not in st.session_state:
+    st.session_state.message_history =  []
+history = st.session_state.message_history
+
+# print(st.session_state.message_history)
+for i, message_ in enumerate(st.session_state.message_history):
+    message(message_,is_user=i % 2 == 0, key=i) # display all the previous message
+
+# placeholder = st.empty() # placeholder for latest message
+input_ = st.text_input("아무 말이나 해보세요", value="")
+
+if input_ is not None and len(input_) > 0:
+    if len(history) <= 1 or history[-2] != input_:
+        with st.spinner("대답을 생성중입니다..."):
+            st.session_state.message_history.append(input_)
+            response = get_response(tokenizer, model, history)
+            st.session_state.message_history.append(response)
+            st.experimental_rerun()

requirements.txt

@@ -0,0 +1,3 @@
+transformers
+streamlit_chat
+torch