Spaces:
Running
Running
File size: 3,225 Bytes
0f93e9d c05e788 7ac0faf 0f93e9d 8ef4162 0f93e9d c05e788 0f93e9d c05e788 7ac0faf c05e788 7ac0faf 8ef4162 c05e788 8ef4162 7ac0faf c05e788 0f93e9d c05e788 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
import os
import tempfile
import shutil
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from src.knowledge_base.loader import load_documents
from config.settings import VECTOR_STORE_PATH, EMBEDDING_MODEL, HF_TOKEN
from config.constants import CHUNK_SIZE, CHUNK_OVERLAP
def get_embeddings():
"""Получение модели эмбеддингов"""
return HuggingFaceEmbeddings(
model_name=EMBEDDING_MODEL,
model_kwargs={'device': 'cpu'}
)
def create_vector_store():
"""Создание векторного хранилища и загрузка в датасет"""
# Загрузка документов
documents = load_documents()
if not documents:
return False, "Ошибка: документы не загружены"
# Разделение на чанки
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=CHUNK_SIZE,
chunk_overlap=CHUNK_OVERLAP
)
chunks = text_splitter.split_documents(documents)
# Инициализация эмбеддингов
embeddings = get_embeddings()
# Создание векторного хранилища во временной директории
with tempfile.TemporaryDirectory() as temp_dir:
vector_store = FAISS.from_documents(chunks, embeddings)
# Сохраняем во временную директорию
vector_store.save_local(folder_path=temp_dir)
# Копируем файлы в VECTOR_STORE_PATH для последующей загрузки
os.makedirs(VECTOR_STORE_PATH, exist_ok=True)
for file in ["index.faiss", "index.pkl"]:
shutil.copy2(
os.path.join(temp_dir, file),
os.path.join(VECTOR_STORE_PATH, file)
)
# Загрузка в датасет с явной передачей токена
from src.knowledge_base.dataset import DatasetManager
dataset = DatasetManager(token=HF_TOKEN)
success, message = dataset.upload_vector_store()
# Очищаем локальные файлы после загрузки
shutil.rmtree(VECTOR_STORE_PATH)
if not success:
return False, f"Ошибка загрузки в датасет: {message}"
return True, f"База знаний создана успешно! Загружено {len(documents)} документов, создано {len(chunks)} чанков."
def load_vector_store():
"""Загрузка векторного хранилища"""
embeddings = get_embeddings()
if not os.path.exists(os.path.join(VECTOR_STORE_PATH, "index.faiss")):
return None
try:
vector_store = FAISS.load_local(
VECTOR_STORE_PATH,
embeddings,
allow_dangerous_deserialization=True
)
return vector_store
except Exception as e:
print(f"Ошибка загрузки векторного хранилища: {str(e)}")
return None
|