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

# 加载模型和tokenizer
model_name = "您的用户名/text-style-converter"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto"
)

def convert_text_style(input_text):
    """文本风格转换函数"""
    if not input_text.strip():
        return "请输入要转换的文本"
    
    prompt = f"""以下是一个文本风格转换任务，请将书面化、技术性的输入文本转换为自然、口语化的表达方式。

### 输入文本：
{input_text}

### 输出文本：
"""
    
    inputs = tokenizer(prompt, return_tensors="pt")
    
    with torch.no_grad():
        outputs = model.generate(
            inputs.input_ids,
            attention_mask=inputs.attention_mask,
            max_new_tokens=500,
            temperature=0.7,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id
        )
    
    full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # 提取生成的部分
    if "### 输出文本：" in full_response:
        response = full_response.split("### 输出文本：")[-1].strip()
    else:
        response = full_response
    
    return response

# 创建Gradio接口
iface = gr.Interface(
    fn=convert_text_style,
    inputs=gr.Textbox(
        label="输入文本",
        placeholder="请输入需要转换为口语化的书面文本...",
        lines=5
    ),
    outputs=gr.Textbox(
        label="输出文本",
        lines=5
    ),
    title="中文文本风格转换API",
    description="将书面化、技术性文本转换为自然、口语化表达",
    examples=[
        ["乙醇的检测方法包括以下几项: 1. 酸碱度检查:取20ml乙醇加20ml水,加2滴酚酞指示剂应无色。"],
        ["本品为薄膜衣片,除去包衣后显橙红色至暗红色。"]
    ]
)

# 启动应用
if __name__ == "__main__":
    iface.launch()