Update app.py

app.py

@@ -2,12 +2,8 @@ import streamlit as st
 from langchain_community.llms import HuggingFaceHub
 from langchain_community.embeddings import SentenceTransformerEmbeddings
 from langchain_community.vectorstores import FAISS
-import numpy as np
 
-# gemma = 'google/gemma-7b-it';
-gemma = 'google/recurrentgemma-2b-it';
-
-# 2. 准备知识库数据 (示例)
+# 1. 准备知识库数据 (示例)
 knowledge_base = [
     "Gemma 是 Google 开发的大型语言模型。",
     "Gemma 具有强大的自然语言处理能力。",
@@ -16,7 +12,7 @@ knowledge_base = [
     "Gemma 支持多种语言。"
 ]
 
-# 3. 构建向量数据库 (如果需要，仅构建一次)
+# 2. 构建向量数据库 (如果需要，仅构建一次)
 try:
     embeddings = SentenceTransformerEmbeddings(model_name="all-mpnet-base-v2")
     db = FAISS.from_texts(knowledge_base, embeddings)
@@ -24,24 +20,27 @@ except Exception as e:
     st.error(f"向量数据库构建失败：{e}")
     st.stop()
 
-# 4. 问答函数
+# 3. 问答函数
 def answer_question(gemma, temperature, max_length, question):
-    # 1. 初始化 Gemma 模型
+    # 4. 初始化 Gemma 模型
     try:
         llm = HuggingFaceHub(repo_id=gemma, model_kwargs={"temperature": temperature, "max_length": max_length})
     except Exception as e:
         st.error(f"Gemma 模型加载失败：{e}")
         st.stop()
+
+    # 5. 获取答案
     try:
         question_embedding = embeddings.embed_query(question)
-        question_embedding_np = " ".join(map(str, question_embedding))
-        docs_and_scores = db.similarity_search_with_score(question_embedding_np)
-        # 正确处理 docs_and_scores 列表
-        context = "\n".join([doc.page_content for doc, _ in docs_and_scores])  # 使用 join() 方法
-        print(context)
+        question_embedding_str = " ".join(map(str, question_embedding))
+        print('question_embedding: ' + question_embedding_str)
+        docs_and_scores = db.similarity_search_with_score(question_embedding_str)
+
+        context = "\n".join([doc.page_content for doc, _ in docs_and_scores])
+        print('context: ' + context)
 
         prompt = f"请根据以下知识库回答问题：\n{context}\n问题：{question}"
-        print(prompt)
+        print('prompt: ' + prompt)
 
         answer = llm(prompt)
         return answer
@@ -53,8 +52,8 @@ def answer_question(gemma, temperature, max_length, question):
 st.title("Gemma 知识库问答系统")
 
 gemma = st.selectbox("模型", ("google/gemma-2-9b-it", "google/gemma-2-2b-it", "google/recurrentgemma-2b-it"), 2)
-temperature = st.text_area("temperature", "1.0")
-max_length = st.text_area("max_length", "1024")
+temperature = st.text_input("temperature", "1.0")
+max_length = st.text_input("max_length", "1024")
 question = st.text_area("请输入问题", "Gemma 有哪些特点？")
 
 if st.button("提交"):
@@ -62,6 +61,6 @@ if st.button("提交"):
         st.warning("请输入问题！")
     else:
         with st.spinner("正在查询..."):
-            answer = answer_question(gemma, temperature, max_length, question)
+            answer = answer_question(gemma, float(temperature), int(max_length), question)
             st.write("答案：")
             st.write(answer)