간단데모

app.py

@@ -0,0 +1,48 @@
+import streamlit as st
+
+@st.cache(allow_output_mutation=True)
+def get_pipe():
+    from transformers import AutoTokenizer, AutoModelForCausalLM
+    model_name = "heegyu/koalpaca-355m"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    tokenizer.truncation_side = "right"
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    return model, tokenizer
+
+def get_response(tokenizer, model, context):
+    context = f"<usr>{context}\n<sys>"
+    inputs = tokenizer(
+        context, 
+        truncation=True,
+        max_length=512,
+        return_tensors="pt")
+    
+    generation_args = dict(
+        max_length=256,
+        min_length=64,
+        eos_token_id=2,
+        do_sample=True,
+        top_p=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>", "")
+
+    return response
+
+st.title("KoAlpaca-355M")
+
+with st.spinner("loading model..."):
+    model, tokenizer = get_pipe()
+
+input_ = st.text_area("질문해보세요", value="미국과 중국의 갈등의 원인이 뭐야?")
+ok = st.button("물어보기")
+if input_ is not None and ok and len(input_) > 0:
+    with st.spinner("잠시만요"):
+        response = get_response(tokenizer, model, input_)
+        st.text("대답")
+        st.success(response)

requirements.txt

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