Update app.py

app.py

@@ -1,16 +1,17 @@
 import requests
 import streamlit as st
 from streamlit_lottie import st_lottie
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
 
 st.set_page_config(
     page_title="노래 가사 n행시",
     page_icon="💌",
+    layout="wide"
 )
 
 ### Model
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
 tokenizer = AutoTokenizer.from_pretrained("wumusill/final_20man")
 
 @st.cache
@@ -20,6 +21,17 @@ def load_model():
 
 model = load_model()
 
+class poem(object):
+    def __init__(self,letter):
+        self.letter = letter
+
+    def __str__(self):
+        return self.letter
+
+    def __repr__(self):
+        return "'"+self.letter+"'"
+
+@st.cache(show_spinner=False)
 def mind(input_letter):
     # 두음 법칙 사전
     dooeum = {"라":"나", "락":"낙", "란":"난", "랄":"날", "람":"남", "랍":"납", "랑":"낭", 
@@ -96,12 +108,12 @@ def mind(input_letter):
 
         # print(res_l)
 
-    dictionary = {}
+    poem_dict = {}
 
     for letter, res in zip(input_letter, res_l):
-        dictionary[letter] = res
+        poem_dict[poem(letter)] = res
 
-    return dictionary
+    return poem_dict
 
 ###
 
@@ -152,24 +164,13 @@ with row1_1:
 st.write('---')
 
 # Model & Input
-row2_spacer1, row2_1, row2_spacer2, row2_2, row2_spacer3 = st.columns((0.01, 1.5, 0.05, 1.5, 0.01))
-
-# Genre Selector
-if "genre" not in st.session_state:
-    st.session_state.genre = "전체"
-
-with row2_1:
-    st.radio(
-        "장르를 선택해주세요. 👉",
-        key="genre",
-        options=["전체", "발라드", "록/메탈", "힙합", "트로트"],
-    )
+row2_spacer1, row2_1, row2_spacer2 = st.columns((0.01, 1.5, 0.05))
 
 # Word Input
 if "generate" not in st.session_state:
     st.session_state.generate = False
 
-with row2_2:
+with row2_1:
     word_input = st.text_input(
             "n행시에 사용할 단어를 적고 Enter를 눌러주세요. 👇",
             placeholder='한글 단어'
@@ -179,9 +180,8 @@ with row2_2:
         st.write("n행시 단어 :  ", word_input)
 
     if st.button('n행시 제작하기'):
-        result = mind(word_input)
+        with st.spinner('잠시 기다려주세요...'):
+            result = mind(word_input)
+        st.success('완료됐습니다!')
         for r in result:
-            st.write(f'{r} : {result[r]}')
-
-
-
+            st.write(f'{r} : {result[r]}')