coding

app.py

@@ -15,13 +15,13 @@ if "data_list" not in st.session_state:
 if not st.session_state.data_list:
     try:
         with st.spinner("正在读取数据库..."):
-            st.session_state.converter = OpenCC('tw2s') # 'tw2s.json' 表示繁体中文到简体中文的转换
+            converter = OpenCC('tw2s') # 'tw2s.json' 表示繁体中文到简体中文的转换
             dataset = load_dataset("rorubyy/attack_on_titan_wiki_chinese")
             data_list = []
             answer_list = []
             for example in dataset["train"]:
-                converted_answer = st.session_state.converter.convert(example["Answer"])
-                converted_question = st.session_state.converter.convert(example["Question"])
+                converted_answer = converter.convert(example["Answer"])
+                converted_question = converter.convert(example["Question"])
                 answer_list.append(converted_answer)
                 data_list.append({"Question": converted_question, "Answer": converted_answer})
             st.session_state.answer_list = answer_list
@@ -63,7 +63,7 @@ def answer_question(repo_id, temperature, max_length, question):
     if repo_id != st.session_state.repo_id or temperature != st.session_state.temperature or max_length != st.session_state.max_length:
         try:
             with st.spinner("正在初始化 Gemma 模型..."):
-                llm = HuggingFaceHub(repo_id=repo_id, model_kwargs={"temperature": temperature, "max_length": max_length})
+                st.session_state.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
@@ -91,7 +91,7 @@ def answer_question(repo_id, temperature, max_length, question):
             print("本地数据集筛选完成！")
 
         with st.spinner("正在生成答案..."):
-            answer = llm.invoke(prompt)
+            answer = st.session_state.llm.invoke(prompt)
             # 去掉 prompt 的内容
             answer = answer.replace(prompt, "").strip()
             st.success("答案已经生成！")
@@ -113,6 +113,13 @@ with col2:
 
 st.divider()
 
+def generate_answer(repo_id, temperature, max_length, question):
+    result = answer_question(repo_id, float(temperature), int(max_length), question)
+    print('prompt: ' + result["prompt"])
+    print('answer: ' + result["answer"])
+    st.write("生成答案：")
+    st.write(result["answer"])
+
 col3, col4 = st.columns(2)
 with col3:
     if st.button("使用原数据集中的随机问题"):
@@ -120,9 +127,7 @@ with col3:
         random_index = random.randint(0, dataset_size - 1)
         # 读取随机问题
         random_question = st.session_state.data_list[random_index]["Question"]
-        random_question = st.session_state.converter.convert(random_question)
         origin_answer = st.session_state.data_list[random_index]["Answer"]
-        origin_answer = st.session_state.converter.convert(origin_answer)
         print('[]' + str(random_index) + '/' + str(dataset_size) + ']random_question: ' + random_question)
         print('origin_answer: ' + origin_answer)
 
@@ -130,20 +135,12 @@ with col3:
         st.write(random_question)
         st.write("原始答案：")
         st.write(origin_answer)
-        result = answer_question(gemma, float(temperature), int(max_length), random_question)
-        print('prompt: ' + result["prompt"])
-        print('answer: ' + result["answer"])
-        st.write("生成答案：")
-        st.write(result["answer"])
+        generate_answer(gemma, float(temperature), int(max_length), random_question)
 
 with col4:
-    question = st.text_area("请输入问题", "Gemma 有哪些特点？")
+    question = st.text_area("请输入问题", "《进击的巨人》中都有哪些主要角色？")
     if st.button("提交输入的问题"):
         if not question:
             st.warning("请输入问题！")
         else:
-            result = answer_question(gemma, float(temperature), int(max_length), question)
-            print('prompt: ' + result["prompt"])
-            print('answer: ' + result["answer"])
-            st.write("生成答案：")
-            st.write(result["answer"])
+            generate_answer(gemma, float(temperature), int(max_length), question)