88 lines
3.1 KiB
Python
88 lines
3.1 KiB
Python
#novel_generator/knowledge.py
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
知识文件导入至向量库（advanced_split_content、import_knowledge_file）
|
||
"""
|
||
import os
|
||
import logging
|
||
import re
|
||
import traceback
|
||
import nltk
|
||
import warnings
|
||
from utils import read_file
|
||
from novel_generator.vectorstore_utils import load_vector_store, init_vector_store
|
||
from langchain.docstore.document import Document
|
||
|
||
# 禁用特定的Torch警告
|
||
warnings.filterwarnings('ignore', message='.*Torch was not compiled with flash attention.*')
|
||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||
|
||
def advanced_split_content(content: str, similarity_threshold: float = 0.7, max_length: int = 500) -> list:
|
||
"""使用基本分段策略"""
|
||
nltk.download('punkt', quiet=True)
|
||
nltk.download('punkt_tab', quiet=True)
|
||
sentences = nltk.sent_tokenize(content)
|
||
if not sentences:
|
||
return []
|
||
|
||
final_segments = []
|
||
current_segment = []
|
||
current_length = 0
|
||
|
||
for sentence in sentences:
|
||
sentence_length = len(sentence)
|
||
if current_length + sentence_length > max_length:
|
||
if current_segment:
|
||
final_segments.append(" ".join(current_segment))
|
||
current_segment = [sentence]
|
||
current_length = sentence_length
|
||
else:
|
||
current_segment.append(sentence)
|
||
current_length += sentence_length
|
||
|
||
if current_segment:
|
||
final_segments.append(" ".join(current_segment))
|
||
|
||
return final_segments
|
||
|
||
def import_knowledge_file(
|
||
embedding_api_key: str,
|
||
embedding_url: str,
|
||
embedding_interface_format: str,
|
||
embedding_model_name: str,
|
||
file_path: str,
|
||
filepath: str
|
||
):
|
||
logging.info(f"开始导入知识库文件: {file_path}, 接口格式: {embedding_interface_format}, 模型: {embedding_model_name}")
|
||
if not os.path.exists(file_path):
|
||
logging.warning(f"知识库文件不存在: {file_path}")
|
||
return
|
||
content = read_file(file_path)
|
||
if not content.strip():
|
||
logging.warning("知识库文件内容为空。")
|
||
return
|
||
paragraphs = advanced_split_content(content)
|
||
from embedding_adapters import create_embedding_adapter
|
||
embedding_adapter = create_embedding_adapter(
|
||
embedding_interface_format,
|
||
embedding_api_key,
|
||
embedding_url if embedding_url else "http://localhost:11434/api",
|
||
embedding_model_name
|
||
)
|
||
store = load_vector_store(embedding_adapter, filepath)
|
||
if not store:
|
||
logging.info("Vector store does not exist or load failed. Initializing a new one for knowledge import...")
|
||
store = init_vector_store(embedding_adapter, paragraphs, filepath)
|
||
if store:
|
||
logging.info("知识库文件已成功导入至向量库(新初始化)。")
|
||
else:
|
||
logging.warning("知识库导入失败，跳过。")
|
||
else:
|
||
try:
|
||
docs = [Document(page_content=str(p)) for p in paragraphs]
|
||
store.add_documents(docs)
|
||
logging.info("知识库文件已成功导入至向量库(追加模式)。")
|
||
except Exception as e:
|
||
logging.warning(f"知识库导入失败: {e}")
|
||
traceback.print_exc()
|