import gradio as gr
from transformers import pipeline

# モデル読み込み（起動時に一度だけ）
tsundere_ai = pipeline(
    "text-generation",
    model="rinna/japanese-gpt-neox-3.6b-instruction-ppo",
    device_map=None,                  # CPUのみ
    tokenizer_kwargs={"use_fast": False}  # slow tokenizer
)

def convert(text):
    # 入力文が長すぎる場合は切り詰める
    text = text[:200]  # 先頭200文字だけ処理

    prompt = f"""
以下の文章を自然でなめらかなツンデレ口調に変換してください。
・不自然な語尾にしない
・元文の意味を保持する
・一般的な可愛いツンデレの話し方にする

【元の文章】
{text}

【変換後】"""

    # max_new_tokens を短めに
    out = tsundere_ai(prompt, max_new_tokens=50, temperature=0.8)[0]["generated_text"]
    return out.split("【変換後】")[-1].strip()

# Gradio UI
app = gr.Interface(
    fn=convert,
    inputs=gr.Textbox(label="入力文"),
    outputs=gr.Textbox(label="ツンデレ変換結果"),
    title="AIツンデレ変換"
)

app.launch()