import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# 模型名称（与你上传一致）
model_name = "yxccai/text-style-converter"

# 加载 tokenizer 和 model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto"
)

# 构建推理函数
def convert_style(input_text):
    prompt = f"""以下是一个文本风格转换任务，请将书面化、技术性的输入文本转换为自然、口语化的表达方式。

### 输入文本：
{input_text}

### 输出文本：
"""

    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(
        **inputs,
        max_new_tokens=300,
        temperature=0.7,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id
    )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # 提取输出部分
    if "### 输出文本：" in response:
        return response.split("### 输出文本：")[-1].strip()
    return response.strip()

# 构建界面
iface = gr.Interface(
    fn=convert_style,
    inputs=gr.Textbox(label="输入书面化技术性文本"),
    outputs=gr.Textbox(label="口语化转换结果"),
    title="中文文本风格转换器",
    description="将正式、技术性文本转换为自然、口语化表达的语言风格转换模型。",
    examples=[
        ["乙醇的检测方法包括酸碱度检查。"],
        ["本品为薄膜衣片，除去包衣后显橙红色。"]
    ]
)

iface.launch()