集成

answerQuestions.py

@@ -8,30 +8,35 @@ chroma_client = chromadb.PersistentClient(path="data/chroma")
 
 vectordb = chroma_client.get_collection(name="collection")
 
+def generateStory(question):
+    query = [question, ]
+    response = erniebot.Embedding.create(
+        input=query,
+        model='ernie-text-embedding',
+    )
+    query_embeddings = response.data[0]['embedding']
+    print(query_embeddings)
+
+    results = vectordb.query(
+        query_embeddings=query_embeddings,
+        n_results=1  # top_k
+    )
+
+    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
+    )
+    story = response.get_result()
+    return story
+
+
+
 question = "什么是克隆羊"
-query = [question,]
-response = erniebot.Embedding.create(
-    input=query,
-    model='ernie-text-embedding',
-)
-query_embeddings = response.data[0]['embedding']
-print(query_embeddings)
-
-results = vectordb.query(
-    query_embeddings = query_embeddings,
-    n_results=1 #top_k
-)
-
-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
-)
-story = response.get_result()
-print(story)
+print(generateStory(question))

kg.py

@@ -1,6 +1,7 @@
 # encoding: utf-8
 import erniebot
 import json
+from answerQuestions import generateStory
 
 erniebot.api_type = "aistudio"
 erniebot.access_token = "ff1531c8c0f429f92adbc2eaed2e23bfb5349e0f"
@@ -96,17 +97,17 @@ def generate_text2img_prompt(input: str):
     return result
 
 
-story = """
-小朋友，你听说过“克隆”这个词吗？让我来给你讲一个有趣的故事吧！
-
-有一天，科学家们想出了一个非常神奇的方法，可以让一个生命体“复制”出另一个一模一样的生命体，这个方法就叫做“克隆”。你可以把“克隆”想象成是一台神奇的复印机，但是这台复印机不是复印纸张，而是复印生命哦！
-
-1997年的时候，科学家们用这种方法成功地“复印”出了一只名叫多利的绵羊。你知道吗，多利羊并没有爸爸，它是由一只母羊的细胞复制而来的。这就像是我们从电脑上复制一个文件一样，只不过这次我们复制的是一个活生生的羊！
-
-克隆羊多利的诞生，就像是打开了一扇新的大门，让人们看到了生命的奇妙和科学的无限可能。从此以后，人们开始更加关注和研究克隆技术，希望用它来帮助我们解决更多的问题。
-
-所以，克隆羊就是通过克隆技术创造出来的一只羊，它没有父亲，只有母亲，是科学家们用神奇的“生命复印机”复制出来的哦！
-"""
+# story = """
+# 小朋友，你听说过“克隆”这个词吗？让我来给你讲一个有趣的故事吧！
+#
+# 有一天，科学家们想出了一个非常神奇的方法，可以让一个生命体“复制”出另一个一模一样的生命体，这个方法就叫做“克隆”。你可以把“克隆”想象成是一台神奇的复印机，但是这台复印机不是复印纸张，而是复印生命哦！
+#
+# 1997年的时候，科学家们用这种方法成功地“复印”出了一只名叫多利的绵羊。你知道吗，多利羊并没有爸爸，它是由一只母羊的细胞复制而来的。这就像是我们从电脑上复制一个文件一样，只不过这次我们复制的是一个活生生的羊！
+#
+# 克隆羊多利的诞生，就像是打开了一扇新的大门，让人们看到了生命的奇妙和科学的无限可能。从此以后，人们开始更加关注和研究克隆技术，希望用它来帮助我们解决更多的问题。
+#
+# 所以，克隆羊就是通过克隆技术创造出来的一只羊，它没有父亲，只有母亲，是科学家们用神奇的“生命复印机”复制出来的哦！
+# """
 
 
 def judge_img(url: str, kg):
@@ -228,6 +229,8 @@ def story_to_imgs(story: str):
 
 
 if __name__ == "__main__":
+    question = "什么是克隆羊"
+    story = generateStory(question)
 
     res = story_to_imgs(story)
     print(res)