import torch
from transformers import T5ForConditionalGeneration, T5Tokenizer
import gradio as gr

best_model_path = "swcrazyfan/Kingify-2Way-T5-Large-v1_1"
model = T5ForConditionalGeneration.from_pretrained(best_model_path)
tokenizer = T5Tokenizer.from_pretrained("swcrazyfan/Kingify-2Way-T5-Large-v1_1")

def tokenize_data(text, dekingify):
    if dekingify == "Dekingify":
        input_ = "dekingify: " + str(text) + " </s>"
    else:
        input_ = "kingify: " + str(text) + " </s>"
    max_len = 512
    tokenized_inputs = tokenizer(
        input_,
        padding="max_length",
        truncation=True,
        max_length=max_len,
        return_attention_mask=True,
        return_tensors="pt"
    )
    return {"input_ids": tokenized_inputs["input_ids"],
            "attention_mask": tokenized_inputs["attention_mask"]}

def generate_answers(text, dekingify):
    inputs = tokenize_data(text, dekingify)
    results = model.generate(
        input_ids=inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        do_sample=True,
        num_beams=5,
        max_length=512,
        min_length=1,
        early_stopping=True,
        num_return_sequences=1
    )
    answer = tokenizer.decode(results[0], skip_special_tokens=True)
    return answer

iface = gr.Interface(
    title="Kingify 2Way",
    description=("This is a custom AI model that translates modern English into 17th-century English or "
                 "'King James' English (and vice versa). Write anything below, select 'Kingify' or 'Dekingify', "
                 "and click submit."),
    fn=generate_answers,
    inputs=[
        gr.Textbox(label="Original Text", lines=10),
        gr.Radio(label="What do you want to do?", choices=["Kingify", "Dekingify"])
    ],
    outputs="text"
)

iface.launch()