Update app.py

app.py

@@ -7,31 +7,38 @@ import faiss
 from sentence_transformers import SentenceTransformer
 from huggingface_hub import InferenceClient
 
-# ===================== 📚 RAG 模块 =====================
+# ======================================================
+# 📚 RAG 知识库模块
+# ======================================================
 
+KNOWLEDGE_DIR = "knowledge"  # 存放小说txt的文件夹
 EMBED_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
-EMBEDDING_PATH = "embeddings.index"
-CHUNKS_PATH = "text_chunks.json"
+EMBEDDING_PATH = os.path.join(KNOWLEDGE_DIR, "embeddings.index")
+CHUNKS_PATH = os.path.join(KNOWLEDGE_DIR, "text_chunks.json")
+
+def build_knowledge_base():
+    """从 space/knowledge/*.txt 构建知识库（仅首次执行）"""
+    os.makedirs(KNOWLEDGE_DIR, exist_ok=True)
+    txt_files = [os.path.join(KNOWLEDGE_DIR, f) for f in os.listdir(KNOWLEDGE_DIR) if f.endswith(".txt")]
+    if not txt_files:
+        print("⚠️ 未检测到小说文件，请将小说txt放入 space/knowledge/ 文件夹。")
+        return "❌ 未检测到小说文件，请先上传至 knowledge/。"
 
-def load_or_create_embeddings(novel_files):
-    """上传txt文件 -> 切分 -> 生成embedding -> 存储"""
     model = SentenceTransformer(EMBED_MODEL_NAME)
     texts = []
 
-    for file_path in novel_files:
+    for file_path in txt_files:
         with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
             text = f.read()
-
-        # 分段（每段约500字）
+        # 分段
         chunk_size = 500
         chunks = [text[i:i + chunk_size] for i in range(0, len(text), chunk_size)]
         texts.extend(chunks)
 
-    # 生成向量
+    # 生成embedding
     embeddings = model.encode(texts, show_progress_bar=True, normalize_embeddings=True)
     embeddings = np.array(embeddings, dtype=np.float32)
 
-    # 保存索引
     index = faiss.IndexFlatIP(embeddings.shape[1])
     index.add(embeddings)
     faiss.write_index(index, EMBEDDING_PATH)
@@ -41,14 +48,22 @@ def load_or_create_embeddings(novel_files):
 
     return f"✅ 已建立知识库，共 {len(texts)} 段文本。"
 
+def ensure_knowledge_base():
+    """确保知识库存在（Space启动时调用）"""
+    if os.path.exists(EMBEDDING_PATH) and os.path.exists(CHUNKS_PATH):
+        print("✅ 已检测到现有知识库，直接加载。")
+        return "知识库已加载。"
+    else:
+        print("📘 正在首次构建知识库...")
+        return build_knowledge_base()
+
 def search_similar(query, top_k=5):
-    """根据用户输入检索最相关的小说片段"""
-    if not os.path.exists(EMBEDDING_PATH) or not os.path.exists(CHUNKS_PATH):
+    """检索相关小说片段"""
+    if not (os.path.exists(EMBEDDING_PATH) and os.path.exists(CHUNKS_PATH)):
         return ""
 
     model = SentenceTransformer(EMBED_MODEL_NAME)
     index = faiss.read_index(EMBEDDING_PATH)
-
     with open(CHUNKS_PATH, "r", encoding="utf-8") as f:
         texts = json.load(f)
 
@@ -58,8 +73,9 @@ def search_similar(query, top_k=5):
 
     return "\n".join(results)
 
-
-# ===================== 💬 聊天逻辑 =====================
+# ======================================================
+# 💬 聊天逻辑
+# ======================================================
 
 def respond(
     message,
@@ -79,7 +95,7 @@ def respond(
     messages = [{"role": "system", "content": system_message}]
     messages.extend(history)
 
-    # === 🔍 检索相关小说文本 ===
+    # 🔍 RAG 检索
     related_text = search_similar(message, top_k=5)
     if related_text:
         message = f"以下是参考小说片段：\n{related_text}\n\n请模仿上文风格续写或回答：{message}"
@@ -87,7 +103,6 @@ def respond(
     messages.append({"role": "user", "content": message})
 
     response = ""
-
     for chunk in client.chat_completion(
         messages,
         max_tokens=max_tokens,
@@ -99,16 +114,41 @@ def respond(
         token = ""
         if len(choices) and choices[0].delta.content:
             token = choices[0].delta.content
-
-            # ⚙️ 清除异常标记
+            # 过滤异常标记
             token = re.sub(r"<\|im_(start|end)\|>", "", token)
             token = re.sub(r"(?i)<\|assistant\|>", "", token)
-
         response += token
         yield re.sub(r"<\|im_(start|end)\|>", "", response).strip()
 
 
-# ===================== 🧱 Gradio UI =====================
+# ======================================================
+# 🧱 Gradio 界面
+# ======================================================
 
 chatbot = gr.ChatInterface(
-    res
+    respond,
+    type="messages",
+    additional_inputs=[
+        gr.Textbox(value="你是一位文风细腻、善于模仿小说风格的AI写作助手。", label="System message"),
+        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
+        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p"),
+    ],
+)
+
+with gr.Blocks() as demo:
+    gr.Markdown("# 📖 DeepSeek 小说创作助手（内置知识库版）")
+
+    # 启动时自动加载知识库
+    status = ensure_knowledge_base()
+    gr.Markdown(f"### 🧠 {status}")
+
+    with gr.Sidebar():
+        gr.LoginButton()
+        gr.Markdown("小说知识库路径：`space/knowledge/`")
+        gr.Markdown("如需更新内容，请替换该目录内的txt文件并重新部署Space。")
+
+    chatbot.render()
+
+if __name__ == "__main__":
+    demo.launch()