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

model_name = "cameltech/japanese-gpt-1b-PII-masking"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

instruction = "# タスク\n入力文中の個人情報をマスキングせよ\n\n# 入力文\n"

if torch.cuda.is_available():
    model = model.to("cuda")

def preprocess(text):
    return text.replace("\n", "<LB>")

def postprocess(text):
    return text.replace("<LB>", "\n")

generation_config = {
    "max_new_tokens": 256,
    "num_beams": 3,
    "num_return_sequences": 1,
    "early_stopping": True,
    "eos_token_id": tokenizer.eos_token_id,
    "pad_token_id": tokenizer.pad_token_id,
    "repetition_penalty": 3.0
}

def generate(input_text):
    input_text = instruction + input_text
    input_text += tokenizer.eos_token
    input_text = preprocess(input_text)
    
    with torch.no_grad():
        token_ids = tokenizer.encode(input_text, add_special_tokens=False, return_tensors="pt")
    
        output_ids = model.generate(
            token_ids.to(model.device),
            **generation_config
        )
    output = tokenizer.decode(output_ids.tolist()[0][token_ids.size(1) :], skip_special_tokens=True)
    return postprocess(output)


iface = gr.Interface(
    fn=generate,
    inputs=gr.Textbox(label = "input", default = "渡邉亮です。現在の住所は東京都世田谷区代沢1-2-3です。"),
    outputs="text",
    title = "Japanese GPT 1B PII Masking DEMO"
)

iface.launch()