import streamlit as st
from transformers import pipeline

# transformers パイプラインのインポート
fugu_translator_enja = pipeline('translation', model='staka/fugumt-en-ja')
fugu_translator_jaen = pipeline('translation', model='staka/fugumt-ja-en')
zhja_translator = pipeline("translation", model="Helsinki-NLP/opus-mt-tc-big-zh-ja")
jazh_translator = pipeline(model="larryvrh/mt5-translation-ja_zh")

# Streamlit アプリケーション
st.title("Multi-Language Translator")

# st.session_state で session-specific state を作成
if 'session_models' not in st.session_state:
    st.session_state.session_models = {
        'enja': fugu_translator_enja,
        'jaen': fugu_translator_jaen,
        'zhja': zhja_translator,
        'jazh': jazh_translator
    }

# デフォルトの入力値
default_model = 'enja'
default_text = ''

# ユーザー入力の取得
model = st.selectbox("モデル", ['enja', 'jaen', 'zhja', 'jazh'], index=0, key='model')
text = st.text_area("入力テキスト", default_text)

# 翻訳ボタンが押されたときの処理
if st.button("翻訳する"):
    result = st.session_state.session_models[model](text)[0]['translation_text']
    
    # Outputをcollumまたはcontainerに格納
    output_col, _ = st.columns(2)
    output_col.write(f"翻訳結果: {result}")

    # Experimental rerun without re-executing the entire app
    st.experimental_rerun([output_col])