from transformers import pipeline
import gradio as gr

# Pipeline
# Translate from Japanese to English
# translator_ja_to_en = pipeline("translation_ja_to_en", model="Helsinki-NLP/opus-mt-ja-en")
translator_ja_to_en = pipeline("translation_ja_to_en", model="japanese-denim/nllb-finetuned-naga-to-eng")

# Summerize from English to English
# summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
summarizer = pipeline("summarization", model="google/pegasus-large")

# Translate from English to Japanese
# translator_en_to_ja = pipeline("translation_en_to_ja", model="Helsinki-NLP/opus-tatoeba-en-ja")
translator_en_to_ja = pipeline("translation_en_to_ja", model="ZenXir/marian-finetuned-kde4-en-to-ja")


maxlength = 100

def summarize_and_translate(text):
     # check max text length
    if len(text) > maxlength:
        return f"エラー：{maxlength}文字以内で入力してください\nError: The text cannot exceed {maxlength} characters. Please shorten your text."
    
    #  Translate from Japanese to English
    translated_text_to_en = translator_ja_to_en(text, max_length=maxlength+50)[0]['translation_text']

    # Summerize from English to English
    summary_in_en = summarizer(translated_text_to_en, min_length=5, max_length=50)[0]['summary_text']

    #  Translate from English to Japanese
    summary_in_ja = translator_en_to_ja(summary_in_en, max_length=200)[0]['translation_text']

    return summary_in_ja

demo = gr.Interface(
    fn=summarize_and_translate,
    inputs=gr.Textbox(lines=5, placeholder= f"{maxlength}文字以内で日本語の文章を入力してください\nPlease enter Japanese text within a {maxlength}-character limit"),
    outputs=gr.Textbox(lines=5, label="日本語要約 (Summarized Japanese)"),
    title="日本語要約 Jpanese Summarizer",
    description=f"日本語の文章を入力すると、日本語の要約が表示されます。(Enter Japanese text up to {maxlength} characters to see its summary)"
)

demo.launch()