app.py

import os
import pandas as pd
import gradio as gr
from openai import OpenAI

# 初始化 OpenAI 客户端
client = OpenAI(
    api_key=os.getenv("DEEPSEEK_API_KEY"),
    base_url="https://api.deepseek.com"
)

def read_csv_prefix(file, n_lines=128):
    """读取 CSV 文件的前 n 行并提取统一前缀"""
    df = pd.read_csv(file.name)
    if 'prompt' in df.columns:
        return df['prompt'].head(n_lines).tolist()
    else:
        raise ValueError("CSV 文件中没有 'prompt' 列")

def generate_prompt(user_input, csv_prompts, language='English'):
    """基于用户输入和 CSV 提示生成新提示词"""
    # 统一前缀
    prefix = "In the style of anime landscape,"
    
    # 构建系统提示
    system_prompt = (
        f"你是一位自然诗人，请将以下内容转化为{'英文' if language == 'English' else '中文'}的风景描写。"
        "请确保风格与提供的 CSV 中的描述一致，并以统一前缀 '{prefix}' 开头。"
    )
    
    # 拼接用户输入和 CSV 内容
    user_content = f"转换以下内容：\n{user_input}\n\n参考示例：\n" + "\n".join(csv_prompts[:5])
    
    response = client.chat.completions.create(
        model="deepseek-chat",
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_content}
        ],
        temperature=0.7,
        max_tokens=300
    )
    return response.choices[0].message.content.strip()

# Gradio 界面定义
with gr.Blocks(title="提示词生成器", theme=gr.themes.Soft()) as app:
    gr.Markdown("## 🎨 基于 CSV 的提示词生成器")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="输入文本",
                placeholder="例如：森林里的阳光透过树叶洒落...",
                lines=5
            )
            csv_file = gr.File(label=".csv 文件", file_types=[".csv"])
            line_count = gr.Number(label="读取行数", value=128, minimum=1, maximum=10000)
            language_choice = gr.Radio(choices=["English", "Chinese"], label="输出语言", value="English")
            btn = gr.Button("生成提示词", variant="primary")
        with gr.Column():
            output_text = gr.Textbox(
                label="生成的提示词",
                lines=8,
                interactive=False
            )
    
    # 示例数据
    examples = gr.Examples(
        examples=[
            ["清晨的阳光洒在湖面上"],
            ["夜晚的城市灯火璀璨"]
        ],
        inputs=[input_text],
        label="点击试试示例"
    )

    # 事件绑定
    btn.click(
        fn=lambda text, file, lines, lang: generate_prompt(text, read_csv_prefix(file, lines), lang),
        inputs=[input_text, csv_file, line_count, language_choice],
        outputs=output_text
    )

app.launch(share=True)