import streamlit as st
import pandas as pd
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'], index=0, key='model')
text = st.text_area("入力テキスト", default_text)

# 翻訳ボタンが押されたときの処理
if st.button("翻訳する"):
    # Perform translation
    result = st.session_state.session_models[model](text)[0]['translation_text']
    
    # Save the data to a CSV file
    data = {'ID': [1], 'Original Text': [text], 'Result': [result]}
    df = pd.DataFrame(data)
    df.to_csv('translation_data.csv', mode='a', header=not st.session_state.csv_created, index=False)
    
    # Update the CSV creation flag
    st.session_state.csv_created = True
    
    # Display the result
    st.write(f"翻訳結果: {result}")