import gradio as gr
from transformers import AutoTokenizer

# トークナイザーの読み込み
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-72B-Instruct")

# エンコードとデコードを行う関数
def process_text(text, tokens):
    if text:
        inputs = tokenizer(text, return_tensors="pt")
        encoded_tokens = inputs['input_ids'].squeeze().tolist()
        decoded_text = tokenizer.decode(encoded_tokens)
        return encoded_tokens, decoded_text
    elif tokens:
        # トークン列をリストに変換
        tokens_list = [int(token) for token in tokens.split(',')]
        decoded_text = tokenizer.decode(tokens_list)
        return tokens_list, decoded_text
    return [], ""

# Gradioインターフェースの定義
iface = gr.Interface(
    fn=process_text,
    inputs=[
        gr.Textbox(label="エンコード用テキスト"),
        gr.Textbox(label="デコード用トークン列 (カンマ区切り)")
    ],
    outputs=[
        gr.JSON(label="エンコードされたトークン"),
        gr.Textbox(label="デコードされたテキスト")
    ],
    title="トークナイザーインターフェース",
    description="テキストをエンコード・デコードします。"
)

# インターフェースの起動
iface.launch()