import torch
import gradio as gr

from unsloth import FastLanguageModel
from peft import PeftModel

# =========================
# Load model once at startup
# =========================

print("Loading base model...")

base_model, proc = FastLanguageModel.from_pretrained(
    "unsloth/Qwen3.5-9B",
    max_seq_length=2048,
    load_in_4bit=True,  # Recommended unless you have lots of VRAM
)

tokenizer = proc.tokenizer if hasattr(proc, "tokenizer") else proc

print("Loading LoRA adapter...")

model = PeftModel.from_pretrained(
    base_model,
    "XiangJinYu/Qwen3.5-9B-Humanize-DPO-Round2",
    is_trainable=False,
)

if hasattr(model, "config") and getattr(model.config, "model_type", "") == "qwen3_5":
    model.config.model_type = "qwen3"

FastLanguageModel.for_inference(model)

print("Model loaded successfully!")

# =========================
# Inference function
# =========================

def humanize_text(
    text,
    temperature,
    top_p,
    max_tokens,
):
    if not text.strip():
        return ""

    instruction = (
        "请将下面文本改写得更像自然人写作，"
        "保持原意与事实，不要加标题或说明。"
    )

    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": f"{instruction}\n\n原文：{text}",
                }
            ],
        }
    ]

    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,
    )

    inputs = tokenizer(
        prompt,
        return_tensors="pt",
    ).to(model.device)

    with torch.inference_mode():
        outputs = model.generate(
            **inputs,
            max_new_tokens=int(max_tokens),
            temperature=float(temperature),
            top_p=float(top_p),
            do_sample=True,
            repetition_penalty=1.1,
        )

    generated = outputs[0][inputs["input_ids"].shape[1]:]

    result = tokenizer.decode(
        generated,
        skip_special_tokens=True,
    )

    return result.strip()


# =========================
# Gradio UI
# =========================

with gr.Blocks(title="Qwen Humanizer") as demo:
    gr.Markdown(
        """
        # Qwen Humanizer

        Paste academic, AI-generated, or formal text and rewrite it to sound more natural while preserving meaning.
        """
    )

    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="Input Text",
                lines=12,
                placeholder="Paste text here...",
            )

            temperature = gr.Slider(
                minimum=0.1,
                maximum=1.2,
                value=0.65,
                step=0.05,
                label="Temperature",
            )

            top_p = gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.9,
                step=0.05,
                label="Top P",
            )

            max_tokens = gr.Slider(
                minimum=64,
                maximum=1024,
                value=512,
                step=32,
                label="Max New Tokens",
            )

            btn = gr.Button("Humanize")

        with gr.Column():
            output_text = gr.Textbox(
                label="Humanized Output",
                lines=12,
            )

    btn.click(
        fn=humanize_text,
        inputs=[
            input_text,
            temperature,
            top_p,
            max_tokens,
        ],
        outputs=output_text,
    )

demo.launch()