from transformers import pipeline, MarianMTModel, MarianTokenizer
import gradio as gr

# Load the summarization model
summarization_model = pipeline("summarization", model="dewifaj/summarizer_samsum_model")

# Define the model and the tokenizer
model_name = 'Helsinki-NLP/opus-mt-en-id'
tokenizer = MarianTokenizer.from_pretrained(model_name)
model = MarianMTModel.from_pretrained(model_name)


def summarize(text):
    # Summarize the input text
    summary = summarization_model(text, max_length=300, min_length=30, do_sample=False)[0]['summary_text']
    return summary

def translate_text(summary):
    # Encode the text
    input_ids = tokenizer.encode(summary, return_tensors='pt')
    # Generate translation
    output = model.generate(input_ids)
    translation = tokenizer.decode(output[0], skip_special_tokens=True)
    return translation

def summarize_and_translate(text):
    summarized_text = summarize(text)
    translated_text = translate_text(summarized_text)
    return translated_text

example_conversations = [
    "Laras: Hey Bayu, did you finish the math homework that Ms. Smith assigned?\nBayu: Yeah, I managed to get through most of it. How about you?\nLaras: I'm struggling a bit with the last few questions. The ones involving quadratic equations always trip me up.\nBayu: Ah, I see. Those can be tricky. Do you want some help with them?\nLaras: That would be awesome, thanks! I just can't seem to wrap my head around factoring them properly.\nBayu: No problem at all. Let's take a look. So, for this equation here, we need to find two numbers that multiply to give us the constant term and add up to give us the coefficient of the middle term.\nLaras: Right, I remember that part. But sometimes finding those numbers feels like looking for a needle in a haystack.\nBayu: I get what you mean. It takes some practice to spot the patterns. Let's break it down step by step. First, we write down the equation...\n(They start working on the problem together, going through the steps of factoring and solving the quadratic equation.)\nLaras: Ah, I think I'm starting to get it now. It's all about finding the right pair of numbers.\nBayu: Exactly! Once you get the hang of it, factoring becomes much easier. And remember, practice makes perfect.\nLaras: Thanks a lot, Bayu. I really appreciate your help.\nBayu: No problem, Laras. Anytime you need assistance, just let me know. We're in this together!",
    "Louis: What did you do over the weekend?\nCharlie: I went hiking with some friends. It was a lot of fun!\nLouis: That sounds like a great way to spend the weekend. Where did you go hiking?\nCharlie: We went to the nearby national park. The trails were beautiful this time of year.\nLouis: I've been meaning to go there. Maybe next weekend!",
    "Bill: Have you seen the latest movie?\nAnne: Yes, I watched it last night. It was amazing!\nBill: I heard it got mixed reviews. What did you like about it?\nAnne: The cinematography was stunning, and the storyline kept me hooked until the end of the movie.\nBill: Interesting! Maybe I'll watch it this weekend. Thanks for sharing your thoughts."
]

# Create a Gradio interface
gr.Interface(
    fn=summarize_and_translate,
    inputs=gr.Textbox(lines=10, label="Masukkan Percakapan dalam Bahasa Inggris"),
    outputs=gr.Textbox(lines=5, label="Ringkasan Teks (Bahasa Indonesia)"),
    title="ConvoCatch: English to Bahasa Indonesia Conversational Digest",
    description="Ubah percakapan panjang dengan mudah! ConvoCatch adalah alat yang tepat untuk merangkum percakapan dari Bahasa Inggris ke Bahasa Indonesia dengan mudah.",
    examples=example_conversations
).launch()