整合后前端跑通流程

answerQuestions.py

@@ -2,41 +2,42 @@ import pandas as pd
 import chromadb
 import os
 import erniebot
+
+# 需要设置token
 erniebot.api_type = "aistudio"
-erniebot.access_token = "ff1531c8c0f429f92adbc2eaed2e23bfb5349e0f"
+erniebot.access_token = ""
+
 chroma_client = chromadb.PersistentClient(path="data/chroma")
 
 vectordb = chroma_client.get_collection(name="collection")
 
+
 def generateStory(question):
-    query = [question, ]
+    query = [
+        question,
+    ]
     response = erniebot.Embedding.create(
         input=query,
-        model='ernie-text-embedding',
+        model="ernie-text-embedding",
     )
-    query_embeddings = response.data[0]['embedding']
+    query_embeddings = response.data[0]["embedding"]
     print(query_embeddings)
 
-    results = vectordb.query(
-        query_embeddings=query_embeddings,
-        n_results=1  # top_k
-    )
+    results = vectordb.query(query_embeddings=query_embeddings, n_results=1)  # top_k
 
-    search_result = {'metadatas': results['metadatas'][0], 'documents': results['documents'][0]}
+    search_result = {
+        "metadatas": results["metadatas"][0],
+        "documents": results["documents"][0],
+    }
     prompt = f"你是一个擅长给小朋友讲故事的小助手，请根据自己的知识和检索到的内容，对小朋友的疑问进行解答。注意，由于小朋友的年龄较小，问题解答需要用故事的形式。尽量使故事生动、有趣、富有哲理。 \n 用户问题: {question} \n\n 搜索结果:\n {search_result}"
     print(prompt)
 
-    messages = [
-        {"role": "user", "content": prompt}
-    ]
-    response = erniebot.ChatCompletion.create(
-        model="ernie-4.0",
-        messages=messages
-    )
+    messages = [{"role": "user", "content": prompt}]
+    response = erniebot.ChatCompletion.create(model="ernie-4.0", messages=messages)
     story = response.get_result()
     return story
 
 
-
-question = "什么是克隆羊"
-print(generateStory(question))
+if __name__ == "__main__":
+    question = "什么是克隆羊"
+    print(generateStory(question))

img.py

@@ -3,9 +3,8 @@ import dashscope
 import os
 from dashscope import ImageSynthesis
 
-os.environ["DASHSCOPE_API_KEY"] = "sk-9f398377a14c47249ce53ca4c127a67c"
-
-import erniebot
+# 需要配置环境变量
+os.environ["DASHSCOPE_API_KEY"] = "sk-"
 
 
 def judge_image(url: str, text: str):

kg.py

@@ -1,10 +1,11 @@
 # encoding: utf-8
 import erniebot
 import json
-from answerQuestions import generateStory
 
+
+# 需要配置环境变量
 erniebot.api_type = "aistudio"
-erniebot.access_token = "ff1531c8c0f429f92adbc2eaed2e23bfb5349e0f"
+erniebot.access_token = ""
 
 
 def res2json(res: str):
@@ -220,17 +221,29 @@ def story_to_imgs(story: str):
     print("\n绘本生成结果：")
     for item in res:
         print(item["caption"], item["url"])
-        url = item["url"]
-        # 保存到本地
-        file_name = item["caption"] + ".png"
-        with open("images/%s" % file_name, "wb+") as f:
-            f.write(requests.get(url).content)
+
+        # url = item["url"]
+        # # 保存到本地
+        # file_name = item["caption"] + ".png"
+        # with open("images/%s" % file_name, "wb+") as f:
+        #     f.write(requests.get(url).content)
+
     return res
 
 
 if __name__ == "__main__":
     question = "什么是克隆羊"
-    story = generateStory(question)
+    story = """小朋友，你听说过“克隆”这个词吗？让我来给你讲一个有趣的故事吧！
+
+有一天，科学家们想出了一个非常神奇的方法，可以让一个生命体“复制”出另一个一模一样的生命体，这个方法就叫做“克隆”。你可以把“克隆”想象成是一台神奇的复印机，但是这台复印机不是复印纸张，而是复印生命哦！
+
+1997年的时候，科学家们用这种方法成功地“复印”出了一只名叫多利的绵羊。你知道吗，多利羊并没有爸爸，它是由一只母羊的细胞复制而来的。这就像是我们从电脑上复制一个文件一样，只不过这次我们复制的是一个活生生的羊！
+
+克隆羊多利的诞生，就像是打开了一扇新的大门，让人们看到了生命的奇妙和科学的无限可能。从此以后，人们开始更加关注和研究克隆技术，希望用它来帮助我们解决更多的问题。
+
+所以，克隆羊就是通过克隆技术创造出来的一只羊，它没有父亲，只有母亲，是科学家们用神奇的“生命复印机”复制出来的哦！
+"""
+    # story = generateStory(question)
 
     res = story_to_imgs(story)
     print(res)

main.py

@@ -29,9 +29,12 @@ def mock_story(query: str):
 所以，克隆羊就是通过克隆技术创造出来的一只羊，它没有父亲，只有母亲，是科学家们用神奇的“生命复印机”复制出来的哦！"""
 
 
-def generate_image(text):
-    """ """
-    res = mock_img(text)
+from kg import story_to_imgs
+
+
+def generate_image(story):
+
+    res = story_to_imgs(story)
     md = ""
     # 改为markdown
     for i in range(len(res)):
@@ -39,31 +42,43 @@ def generate_image(text):
     return md
 
 
+from answerQuestions import generateStory
+
+
 def generate_story(query):
-    """ """
-    return mock_story(query)
-
-
-def generate_story_and_image(query):
-    """ """
-    story = generate_story(query)
-    image = generate_image(query)
-    return story, image
-
-
-# 生成一个gradio界面
-demo = gr.Interface(
-    fn=generate_story_and_image,
-    inputs=gr.Textbox(lines=5, label="问题"),
-    outputs=[
-        gr.Textbox(lines=15, label="故事"),
-        gr.Markdown(label="绘本"),
-    ],
-    title="绘本生成器",
-    description="输入一个问题，生成一个故事和绘本图片",
-    theme="compact",
-)
 
+    return generateStory(query)
+
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("# 绘本生成器")
+
+    with gr.Row():
+        with gr.Column():
+            input_text = gr.Textbox(label="输入")
+            with gr.Row():
+                clear_btn = gr.Button("清空")
+                story_btn = gr.Button("生成故事", variant="primary")
+            with gr.Row():
+                img_btn = gr.Button("生成绘本", variant="primary")
+        with gr.Column():
+            story_output = gr.Textbox(label="生成的故事", interactive=False, lines=10)
+            story_btn.click(generate_story, inputs=input_text, outputs=story_output)
+        with gr.Column():
+            gr.Markdown("## 生成的绘本")
+            img_output = gr.Markdown(label="生成的绘本")
+            img_btn.click(generate_image, inputs=story_output, outputs=img_output)
+
+        clear_btn.click(
+            lambda x: [
+                gr.update(value=""),
+                gr.update(value=""),
+                gr.update(value=""),
+            ],
+            None,
+            [input_text, story_output, img_output],
+        )
 
 if __name__ == "__main__":
     demo.launch()

readme.md

@@ -1,3 +1,18 @@
+
+# 环境变量配置
+img.py 中 
+os.environ["DASHSCOPE_API_KEY"] = "sk-9f398377a14c47249ce53ca4c127a67c"
+
+
+kg.py 中
+erniebot.api_type = "aistudio"
+erniebot.access_token = "ff1531c8c0f429f92adbc2eaed2e23bfb5349e0f"
+
+
+answerQuestions.py 中
+erniebot.api_type = "aistudio"
+erniebot.access_token = "ff1531c8c0f429f92adbc2eaed2e23bfb5349e0f"
+
 # 图片生成部分
 
 故事：