# -*- coding: utf-8 -*-
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel
import torch

# --- 配置 (Configuration) ---

# 您上传到Hub的仓库ID (基础模型 + LoRA适配器)
# Your repository ID on the Hugging Face Hub (containing the LoRA adapter)
hub_repo_id = "yxccai/text-style"

# ==============================================================================
# !! 性能问题的核心原因与解决方案 !!
# !! CORE PERFORMANCE ISSUE & SOLUTIONS !!
#
# 您遇到的5分钟耗时问题，是因为在免费的Hugging Face Space (CPU环境)上运行了
# 一个1.8B参数的大模型。CPU进行这种规模的计算非常缓慢。
# The 5-minute delay is caused by running a large 1.8B parameter model on a
# free Hugging Face Space, which uses a CPU. CPU inference for a model of this
# size is inherently very slow.
#
# --- 解决方案 (Solutions) ---
#
# 1. (推荐/免费) 使用更小的基础模型 (Recommended/Free):
#    将下面的 `base_model_name` 从 "Qwen/Qwen1.5-1.8B-Chat"
#    更改为 "Qwen/Qwen1.5-0.5B-Chat"。这是最直接有效的免费解决方案。
#    Change `base_model_name` below from "Qwen/Qwen1.5-1.8B-Chat" to
#    "Qwen/Qwen1.5-0.5B-Chat". This is the most effective free solution.
#
# 2. (最佳性能) 升级到GPU硬件 (Best Performance):
#    在Hugging Face Spaces的设置中，将硬件从免费的CPU升级到付费的GPU
#    （例如 T4 small）。这将提供数量级的速度提升。
#    In your Hugging Face Space settings, upgrade the hardware from the free CPU
#    to a paid GPU (e.g., T4 small). This will provide a massive speed-up.
#
# 3. (进阶) 使用量化 (Advanced):
#    加载模型时使用4位或8位量化可以减少内存占用并加速CPU推理，但这会略微
#    牺牲模型精度。这需要更复杂的代码，如下面的 quantization_config 所示。
#    Using 4-bit or 8-bit quantization can reduce memory usage and speed up CPU
#    inference, at a slight cost to model precision. This requires more complex
#    code, like the `quantization_config` shown below.
# ==============================================================================

# --- 应用解决方案1：使用更小的模型 ---
# Applying Solution 1: Use a smaller model
base_model_name = "Qwen/Qwen1.5-0.5B-Chat" # <--- 已修改为0.5B版本以提速

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Gradio App: Using device: {device}")

# --- (可选) 量化配置示例 (Optional) Quantization Config Example ---
# 如果您想尝试方案3，可以取消下面的注释。注意：这在CPU上可能仍然不够快。
# If you want to try Solution 3, you can uncomment the following lines.
# Note: This might still not be fast enough on a CPU.
# bnb_config = BitsAndBytesConfig(
#     load_in_4bit=True,
#     bnb_4bit_use_double_quant=True,
#     bnb_4bit_quant_type="nf4",
#     bnb_4bit_compute_dtype=torch.bfloat16
# )

# --- 加载模型和Tokenizer (Loading Model and Tokenizer) ---
print(f"Gradio App: Loading base model: {base_model_name}")
# 1. 加载基础模型 (Load base model)
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_name,
    torch_dtype="auto",
    device_map="auto", # 使用 "auto" 让transformers库自动处理设备映射
    trust_remote_code=True
    # quantization_config=bnb_config # 如果使用量化，请添加此行
)
# base_model.to(device) # device_map="auto" 时不需要手动 .to(device)

print(f"Gradio App: Loading tokenizer from: {hub_repo_id}")
# 2. 加载Tokenizer (Load Tokenizer)
tokenizer = AutoTokenizer.from_pretrained(hub_repo_id, trust_remote_code=True)
if tokenizer.pad_token is None:
    print("Gradio App: Tokenizer does not have a pad_token, setting it to eos_token.")
    tokenizer.pad_token = tokenizer.eos_token
    base_model.config.pad_token_id = tokenizer.eos_token_id

print(f"Gradio App: Loading LoRA adapter from: {hub_repo_id}")
# 3. 加载并应用LoRA适配器 (Load and apply LoRA adapter)
model = PeftModel.from_pretrained(base_model, hub_repo_id)

model.eval() # 设置为评估模式 (Set to evaluation mode)
print("Gradio App: Model and tokenizer loaded successfully.")

# --- 推理函数 (Inference Function) ---
def chat(input_text):
    if not input_text:
        return "请输入一些文本。"

    print(f"Gradio App: Received input: {input_text}")
    # 构建符合Qwen Chat模板的输入
    # Construct input according to the Qwen Chat template
    messages = [
        {"role": "system", "content": "你是一个文本风格转换助手。请严格按照要求，仅将以下书面文本转换为自然、口语化的简洁表达方式，不要添加任何额外的解释、扩展信息或重复原文。"},
        {"role": "user", "content": input_text}
    ]

    try:
        prompt = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )
    except Exception as e:
        print(f"Error applying chat template: {e}")
        # 如果模板应用失败，提供一个明确的错误信息
        return f"错误：无法应用聊天模板。{e}"

    print(f"Gradio App: Formatted prompt for model:\n{prompt}")
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024).to(device)

    # 生成时禁用梯度计算以节省资源
    # Disable gradient calculations during generation to save resources
    with torch.no_grad():
        generated_ids = model.generate(
            **inputs,
            max_new_tokens=512,  # 减少最大生成长度以加快响应
            do_sample=True,
            temperature=0.7,
            top_p=0.95,
            pad_token_id=tokenizer.eos_token_id
        )

    # 从生成结果中提取回复
    # Extract the reply from the generated result
    response_ids = generated_ids[0][inputs.input_ids.shape[-1]:]
    result = tokenizer.decode(response_ids, skip_special_tokens=True).strip()

    print(f"Gradio App: Extracted result: {result}")
    return result

# --- 创建Gradio界面 (Create Gradio Interface) ---
iface = gr.Interface(
    fn=chat,
    inputs=gr.Textbox(lines=5, label="输入书面文本 (Input Formal Text)"),
    outputs=gr.Textbox(lines=5, label="输出口语化文本 (Output Casual Text)"),
    title="文本风格转换器 (Text Style Converter)",
    description="输入一段书面化的中文文本，模型会尝试将其转换为更自然、口语化的表达方式。由Qwen-0.5B模型微调。(已优化速度)",
    examples=[
        ["乙醇的检测方法包括以下几项: 1. 酸碱度检查:取20ml乙醇加20ml水,加2滴酚酞指示剂应无色,再加1ml 0.01mol/L氢氧化钠应显粉红色."],
        ["本公司今日发布了最新的财务业绩报告，数据显示本季度利润实现了显著增长。"]
    ],
    allow_flagging="never" # 禁用flagging以简化界面
)

if __name__ == "__main__":
    iface.launch()