Update app.py

app.py

@@ -4,11 +4,11 @@ from volcengine.maas import MaasException, ChatRole
 import json
 import os
 
-def maas_chat(input):
+def maas_chat(input, novel_type):  # 添加小说类型参数
     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')
+    maas = MaasService('maas-api.ml-platform-cn-beijing.volces.com', 'cn-beijing',connection_timeout=15)
 
 
     # 使用 os 模块来访问变量和秘密值
@@ -16,21 +16,23 @@ def maas_chat(input):
     sk = os.getenv('SK')
     maas.set_ak(ak)
     maas.set_sk(sk)
-
+    prompt = "我将提供一份小说片段并且是使用 JSON 数组进行格式化，JSON 数组内的每个对象的内容就是一个故事的场景，你根据我提供的场景原文生成用来 AI 绘画的提示词，注意提示词应该和小说片段具有强相关性，当小说内容没有场景时，结合 JSON 中其他内容输出提示词，并同样使用 JSON 输出并且数组大小和输入的数组大小一样，下面给出 JSON 格式化后的文字："
+    if novel_type == "古文":
+        prompt = '我将提供一份古风小说片段并且是使用JSON数组进行格式化，JSON数组内的每个对象的内容就是一个小说内容的一个片段，你根据我提供的小说片段生成中国古代风格的AI绘画的提示词，注意提示词应该和小说片段具有强相关性，当小说内容没有场景时，结合JSON中其他内容输出提示词，并同样使用JSON输出并且数组大小和输入的数组大小一样，下面给出JSON格式化后的文字：';
     # 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_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格式化后的文字："
+                "content": prompt
             }, {
                 "role": ChatRole.USER,
                 "content": query
@@ -39,21 +41,26 @@ def maas_chat(input):
     }
 
     ID = os.getenv('ID')
-    endpoint_id = ID
+    GUID = os.getenv('GUID')
+    # endpoint_id = novel_types =='现代'?ID: GUID
+    endpoint_id = ID if novel_type == "现代" else GUID
     resp = maas.chat(endpoint_id, req)
     output = resp.choices[0].message['content']
-    #去除output'符号
+    #去除 output'符号
 
     list = json.loads(output)
     return '\n'.join(list)
 
 
+novel_types = ["现代", "古文"]  # 定义小说类型列表
+
 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格式）。",
+    inputs=[gr.Textbox(lines=6, placeholder="输入小说片段（JSON 格式）"), gr.Dropdown(novel_types)],  # 添加下拉输入框默认现代
+
+    outputs=gr.Textbox(lines=6, placeholder="输出提示词（JSON 格式）"),
+    title="小说片段转 AI 绘画提示词",
+    description="输入小说片段（JSON 格式），输出用来 AI 绘画的提示词（JSON 格式）。",
 )
 
-demo.launch(share=True)
+demo.launch(share=True)