coding

app.py

@@ -10,13 +10,20 @@ from opencc import OpenCC
 # 原数据集是是繁体中文，为了调试方便，将其转换成简体中文之后使用
 if "dataset_loaded" not in st.session_state:
     st.session_state.dataset_loaded = False
+    st.session_state.data_list = []
+    st.session_state.answer_list = []
 if not st.session_state.dataset_loaded:
     try:
         with st.spinner("正在读取数据库..."):
             converter = OpenCC('tw2s') # 'tw2s.json' 表示繁体中文到简体中文的转换
             dataset = load_dataset("rorubyy/attack_on_titan_wiki_chinese")
-            answer_list = [converter.convert(example["Answer"]) for example in dataset["train"]]
+            for example in dataset["train"]:
+                converted_answer = converter.convert(example["Answer"])
+                converted_question = converter.convert(example["Question"])
+                st.session_state.answer_list.append(converted_answer)
+                st.session_state.data_list.append({"Question": converted_question, "Answer": converted_answer})
             st.success("数据库读取完成！")
+            print("数据库读取完成！")
     except Exception as e:
         st.error(f"读取数据集失败：{e}")
         st.stop()
@@ -29,8 +36,9 @@ if not st.session_state.vector_created:
     try:
         with st.spinner("正在构建向量数据库..."):
             embeddings = SentenceTransformerEmbeddings(model_name="all-mpnet-base-v2")
-            db = FAISS.from_texts(answer_list, embeddings)
+            db = FAISS.from_texts(st.session_state.answer_list, embeddings)
             st.success("向量数据库构建完成！")
+            print("向量数据库构建完成！")
     except Exception as e:
         st.error(f"向量数据库构建失败：{e}")
         st.stop()
@@ -50,6 +58,7 @@ def answer_question(repo_id, temperature, max_length, question):
             with st.spinner("正在初始化 Gemma 模型..."):
                 llm = HuggingFaceHub(repo_id=repo_id, model_kwargs={"temperature": temperature, "max_length": max_length})
                 st.success("Gemma 模型初始化完成！")
+                print("Gemma 模型初始化完成！")
                 st.session_state.repo_id = repo_id
                 st.session_state.temperature = temperature
                 st.session_state.max_length = max_length
@@ -72,12 +81,14 @@ def answer_question(repo_id, temperature, max_length, question):
             print('prompt: ' + prompt)
 
             st.success("本地数据集筛选完成！")
+            print("本地数据集筛选完成！")
 
         with st.spinner("正在生成答案..."):
             answer = llm.invoke(prompt)
             # 去掉 prompt 的内容
             answer = answer.replace(prompt, "").strip()
             st.success("答案已经生成！")
+            print("答案已经生成！")
         return {"prompt": prompt, "answer": answer}
     except Exception as e:
         st.error(f"问答过程出错：{e}")
@@ -98,12 +109,12 @@ st.divider()
 col3, col4 = st.columns(2)
 with col3:
     if st.button("使用原数据集中的随机问题"):
-        dataset_size = len(dataset["train"])
+        dataset_size = len(st.session_state.data_list)
         random_index = random.randint(0, dataset_size - 1)
         # 读取随机问题
-        random_question = dataset["train"][random_index]["Question"]
+        random_question = st.session_state.data_list[random_index]["Question"]
         random_question = converter.convert(random_question)
-        origin_answer = dataset["train"][random_index]["Answer"]
+        origin_answer = st.session_state.data_list[random_index]["Answer"]
         origin_answer = converter.convert(origin_answer)
         print('[]' + str(random_index) + '/' + str(dataset_size) + ']random_question: ' + random_question)
         print('origin_answer: ' + origin_answer)