import gradio as gr

# モデルのロード
model = gr.load("models/Miwa-Keita/zenz-v1-checkpoints")

# 入力を調整する関数
def preprocess_input(user_input):
    prefix = "\uEE00"  # 前に付与する文字列
    suffix = "\uEE01"  # 後ろに付与する文字列
    processed_input = prefix + user_input + suffix
    return model(processed_input)

# 出力を調整する関数
def postprocess_output(model_output):
    suffix = "\uEE01"
    # \uEE01の後の部分を抽出
    if suffix in model_output:
        return model_output.split(suffix)[1]
    return model_output

# インターフェースを定義
iface = gr.Interface(
    fn=lambda x: postprocess_output(preprocess_input(x)),
    inputs=gr.Textbox(label="変換する文字列（カタカナ）"),
    outputs=gr.Textbox(label="変換結果"),
    title="ニューラルかな漢字変換モデルzenz-v1のデモ",
    description="変換したい文字列をカタカナを入力してください"
)

# ローンチ
iface.launch()