Spaces:
Sleeping
Sleeping
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() |