# -*- coding: utf-8 -*-
"""translation.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1bKzjrfpxJqSrJQUuGR27-kRCSndSClBo
"""


from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import gradio as gr

device = "cuda" if torch.cuda.is_available() else "cpu"

language_model_name = "Qwen/Qwen2-1.5B-Instruct"
language_model = AutoModelForCausalLM.from_pretrained(
    language_model_name,
    torch_dtype="auto",
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(language_model_name)

def process_input(input_text, action):
    if action == "Translate to Japanese":
        prompt = f"Please translate the following English text into Japanese: {input_text}"
        lang = "ja"
    elif action == "Translate to English":
        prompt = f"Please translate the following Japanese text into English: {input_text}"
        lang = "en"
    else:
        return "Invalid action. Please choose 'Translate to English' or 'Translate to Japanese'.", "error"

    messages = [
        {"role": "system", "content": "You are a helpful AI assistant for Language Translation."},
        {"role": "user", "content": prompt}
    ]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    model_inputs = tokenizer([text], return_tensors="pt").to(device)

    generated_ids = language_model.generate(
        model_inputs.input_ids,
        max_new_tokens=512
    )
    generated_ids = [
        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
    ]

    output_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
    return output_text, lang

def handle_interaction(input_text, action):
    output_text, lang = process_input(input_text, action)
    return output_text

action_options = ["Translate to English", "Translate to Japanese"]

iface = gr.Interface(
    fn=handle_interaction,
    inputs=[
        gr.Textbox(label="input text"),
        gr.Dropdown(action_options, label="select action")
    ],
    outputs=[
        gr.Textbox(label="output text"),
    ],
    title="Translation App using AI",
    description="Translate input text  based on the selected Language.",
    theme= "gradio/soft"
)

if __name__ == "__main__":
    iface.launch(share=True)