app.py

import gradio as gr
from volcengine.maas.v2 import MaasService
from volcengine.maas import MaasException, ChatRole
import json


def maas_chat(input):
    query = input.split('\n')
    query = [item for item in query if item]
    query = json.dumps(query, ensure_ascii=False)
    maas = MaasService('maas-api.ml-platform-cn-beijing.volces.com', 'cn-beijing')

    # set ak&sk
    maas.set_ak("AKLTYWE0OWIzODc0YjU0NGRjOGE2N2RjMzE2NTUwZjA3OGQ")
    maas.set_sk("T0dJd1pHUXlNemxoWkRJME5HRmlORGxoWWpGa05ESTJOR1EyTURSbFpqTQ==")

    # chat
    req = {
        "parameters": {
            "max_new_tokens": 1000,  # 输出文本的最大tokens限制
            "min_new_tokens": 1,  # 输出文本的最小tokens限制
            "temperature": 0,  # 用于控制生成文本的随机性和创造性，Temperature值越大随机性越大，取值范围0~1
            "top_p": 0,  # 用于控制输出tokens的多样性，TopP值越大输出的tokens类型越丰富，取值范围0~1
            "top_k": 0,  # 选择预测值最大的k个token进行采样，取值范围0-1000，0表示不生效
            "max_prompt_tokens": 4096,  # 最大输入 token 数，如果给出的 prompt 的 token 长度超过此限制，取最后 max_prompt_tokens 个 token 输入模型。
        },
        "messages": [
            {
                "role": ChatRole.SYSTEM,
                "content": "我将提供一份小说片段并且是使用JSON数组进行格式化，JSON数组内的每个对象的内容就是一个故事的场景，你根据我提供的场景原文生成用来AI绘画的提示词，注意提示词应该和小说片段具有强相关性，当小说内容没有场景时，结合JSON中其他内容输出提示词，并同样使用JSON输出并且数组大小和输入的数组大小一样，下面给出JSON格式化后的文字："
            }, {
                "role": ChatRole.USER,
                "content": query
            }
        ]
    }

    endpoint_id = 'ep-20240402173628-s7zd6'
    resp = maas.chat(endpoint_id, req)
    output = resp.choices[0].message['content']
    #去除output'符号

    list = json.loads(output)
    return '\n'.join(list)


demo = gr.Interface(
    fn=maas_chat,
    inputs=gr.Textbox(lines=6, placeholder="输入小说片段（JSON格式）"),
    outputs=gr.Textbox(lines=6, placeholder="输出提示词（JSON格式）"),
    title="小说片段转AI绘画提示词",
    description="输入小说片段（JSON格式），输出用来AI绘画的提示词（JSON格式）。",
)

demo.launch(share=True)