Create temp.py

temp.py

@@ -0,0 +1,138 @@
+import streamlit as st
+import cshogi
+from IPython.display import display
+from transformers import T5ForConditionalGeneration, T5Tokenizer
+import pandas as pd
+
+#モデルの読み込み
+tokenizer = T5Tokenizer.from_pretrained("pizzagatakasugi/shogi_t5", is_fast=True)
+model = T5ForConditionalGeneration.from_pretrained("pizzagatakasugi/shogi_t5_v2")
+model.eval()
+
+st.title("将棋解説文の自動生成")
+df = pd.read_csv("./dataset10.csv")
+num = st.text_input("0から9の数字を入力")
+
+KIFU_TO_SQUARE_NAMES = [
+                    '１一', '１二', '１三', '１四', '１五', '１六', '１七', '１八', '１九',
+                    '２一', '２二', '２三', '２四', '２五', '２六', '２七', '２八', '２九',
+                    '３一', '３二', '３三', '３四', '３五', '３六', '３七', '３八', '３九',
+                    '４一', '４二', '４三', '４四', '４五', '４六', '４七', '４八', '４九',
+                    '５一', '５二', '５三', '５四', '５五', '５六', '５七', '５八', '５九',
+                    '６一', '６二', '６三', '６四', '６五', '６六', '６七', '６八', '６九',
+                    '７一', '７二', '７三', '７四', '７五', '７六', '７七', '７八', '７九',
+                    '８一', '８二', '８三', '８四', '８五', '８六', '８七', '８八', '８九',
+                    '９一', '９二', '９三', '９四', '９五', '９六', '９七', '９八', '９九',
+                ]
+KIFU_FROM_SQUARE_NAMES = [
+                    '11', '12', '13', '14', '15', '16', '17', '18', '19',
+                    '21', '22', '23', '24', '25', '26', '27', '28', '29',
+                    '31', '32', '33', '34', '35', '36', '37', '38', '39',
+                    '41', '42', '43', '44', '45', '46', '47', '48', '49',
+                    '51', '52', '53', '54', '55', '56', '57', '58', '59',
+                    '61', '62', '63', '64', '65', '66', '67', '68', '69',
+                    '71', '72', '73', '74', '75', '76', '77', '78', '79',
+                    '81', '82', '83', '84', '85', '86', '87', '88', '89',
+                    '91', '92', '93', '94', '95', '96', '97', '98', '99',
+                ]
+
+if  num in [str(x) for x in list(range(10))]:
+    df = df.iloc[int(num)]
+    st.write(df["game_type"],df["precedence_name"],df["follower_name"])
+    sfen = df["sfen"].split("\n")
+    bestlist = eval(df["bestlist"])
+    best2list = eval(df["best2list"])
+    te = []
+    te_sf = []
+    movelist = []
+
+    #文字の正規化
+    for x in range(len(sfen)):
+        if x < 2:
+            continue
+        if len(sfen[x]) > 30:
+            te_sf.append(sfen[x])
+        else:
+            #te.append(sfen[x])
+            temp = sfen[x].split()
+            num = temp[1][0] + temp[1][1]
+            for y in range(len(KIFU_FROM_SQUARE_NAMES)):
+                if num == KIFU_FROM_SQUARE_NAMES[y]:
+                    sq = KIFU_TO_SQUARE_NAMES[y]
+            word = sq+temp[1][2:]
+            word = word.replace("竜","龍").replace("成銀","全").replace("成桂","圭").replace("成香","杏")
+            if sfen[x].split()[1] not in ["投了" , "千日手" , "持将棋" , "反則勝ち"]:
+                te.append(temp[0]+" "+word)
+                movelist.append(word)
+            else:
+                movelist.append(sfen[x].split()[1])
+
+    #盤面表示
+    s = st.selectbox(label="手数を選択",options=te)
+
+    with st.expander("parameter"):
+        
+        beams = st.slider("num_beams",min_value=1,max_value=10,step=5,value=1,key=2)
+        tokens = st.slider("min_new_tokens",min_value=0,max_value=50,value=20,key=3)
+        top_p = st.slider("top_p",min_value=0.50,max_value=1.00,value=0.90)
+        top_k = st.slider("top_k",min_value=5,max_value=50,value=30)
+        
+    reload = st.button('盤面生成',key=0)
+    if s in te and reload == True:
+        reload = False
+        idx = te.index(s)
+        board = cshogi.Board(sfen=te_sf[idx+1])
+        st.markdown(board.to_svg(),unsafe_allow_html=True)
+
+        #入力文作成
+        kifs="解説文生成："
+        cnt = 0
+        for kif in movelist:
+            if cnt > idx:
+                break
+            kif = kif.split("(")[0]
+            kifs += kif.replace("▲","").replace("△","")
+            cnt += 1
+            
+        best = ""
+        cnt = 0
+        for x in bestlist[idx]:
+            best += x.split("(")[0]
+            cnt += 1
+            if cnt == 3:
+                break
+            
+        best2 = ""
+        for y in best2list[idx]:
+            best2 += y.split("(")[0]
+            break
+            
+        #st.write(idx,"入力",input)
+        with st.spinner("推論中です..."):
+            input = kifs+"。"+best+"。"+best2
+            tokenized_inputs = tokenizer.encode(
+                                    input, max_length= 512, truncation=True, 
+                                    padding="max_length", return_tensors="pt"
+                                )
+        
+            output_ids = model.generate(input_ids=tokenized_inputs,
+                            max_length=512,
+                            repetition_penalty=10.0, # 同じ文の繰り返しへのペナルティ
+                            do_sample = True,
+                            num_beams = beams,
+                            min_new_tokens = tokens,
+                            top_p = top_p,
+                            top_k = top_k,
+                            num_return_sequences = 5,
+                            )
+            # output_list = []
+            st.write(input)
+            for x in range(5):
+                output_text = tokenizer.decode(output_ids[x], skip_special_tokens=True,
+                                                    clean_up_tokenization_spaces=False)
+                st.write(output_text)
+
+    
+        # temperature = st.slider("temperature",min_value=0.0,max_value=1.0,step=0.01,value=0.3,key=1)
+        # num_beams = st.slider("num_beams",min_value=1,max_value=5,step=1,value=1,key=2)
+        # min_new_tokens = st.slider("min_new_tokens",min_value=0,max_value=100,value=30,key=3)