from transformers import pipeline
import gradio as gr

# Pipeline
# 和訳パイプライン
translator_ja_to_en = pipeline("translation_ja_to_en", model="Helsinki-NLP/opus-mt-ja-en")

# 要約パイプライン(英語)
summarizer = pipeline("summarization", model="facebook/bart-large-cnn")

# 翻訳パイプラインの初期化（英語から日本語）
translator_en_to_ja = pipeline("translation_en_to_ja", model="Helsinki-NLP/opus-tatoeba-en-ja")

def summarize_and_translate(text):
    # 英訳
    translated_text_to_en = translator_ja_to_en(text, max_length=400)[0]['translation_text']

    # 要約
    summary_in_en = summarizer(translated_text_to_en, min_length=5, max_length=20)[0]['summary_text']

    # 和訳
    summary_in_ja = translator_en_to_ja(summary_in_en, max_length=400)[0]['translation_text']

    # 不要なスペースを除去　→　なぜかエラーになるので断念
    # summary_in_ja = re.sub(r'\s+', ' ', summary_in_ja).strip()

    return summary_in_ja

demo = gr.Interface(
    fn=summarize_and_translate,
    inputs=gr.Textbox(lines=5, placeholder="日本語の文章を入力してください"),
    outputs=gr.Textbox(lines=5, label="日本語の要約"),
    title="日本語文書要約ツール",
    description="日本語の文章を入力すると、日本語の要約が表示されます。"
)

demo.launch()