Реализация агента поиска документов с использованием Hugging Face, ChromaDB и Langchain

В современном мире, насыщенном информацией, быстрый поиск релевантных документов имеет решающее значение. Традиционные поисковые системы на основе ключевых слов часто оказываются неэффективными при работе с семантическим смыслом текста. В этой статье мы рассмотрим, как построить мощную поисковую систему для документов, используя:

• Модели встраивания (embedding) от Hugging Face для преобразования текста в векторные представления
• ChromaDB в качестве векторной базы данных для эффективного поиска по сходству
• Sentence Transformers для создания высококачественных текстовых эмбеддингов

Преимущества семантического поиска

Данная реализация обеспечивает возможности семантического поиска – нахождение документов на основе их смысла, а не только совпадения ключевых слов. После выполнения всех шагов вы получите работающую поисковую систему, которая сможет:

• Обрабатывать и встраивать текстовые документы
• Эффективно хранить эти векторные представления
• Находить документы, семантически наиболее близкие к запросу
• Работать с различными типами документов и поисковыми потребностями

Пошаговая реализация

Шаг 1: Установка необходимых библиотек

!pip install chromadb sentence-transformers langchain datasets

Шаг 2: Импорт библиотек

import os
import numpy as np
import pandas as pd
from datasets import load_dataset
import chromadb
from chromadb.utils import embedding_functions
from sentence_transformers import SentenceTransformer
from langchain.text_splitter import RecursiveCharacterTextSplitter
import time

Шаг 3: Загрузка и подготовка данных

Для этого руководства мы будем использовать подмножество статей из Википедии, доступных в библиотеке datasets от Hugging Face:

dataset = load_dataset("wikipedia", "20220301.en", split="train[:1000]")
print(f"Loaded {len(dataset)} Wikipedia articles")


documents = []
for i, article in enumerate(dataset):
   doc = {
       "id": f"doc_{i}",
       "title": article["title"],
       "text": article["text"],
       "url": article["url"]
   }
   documents.append(doc)


df = pd.DataFrame(documents)
df.head(3)

Шаг 4: Разделение документов на фрагменты

Для более детального поиска разделим наши документы на небольшие фрагменты:

text_splitter = RecursiveCharacterTextSplitter(
   chunk_size=1000,
   chunk_overlap=200,
   length_function=len,
)


chunks = []
chunk_ids = []
chunk_sources = []


for i, doc in enumerate(documents):
   doc_chunks = text_splitter.split_text(doc["text"])
   chunks.extend(doc_chunks)
   chunk_ids.extend([f"chunk_{i}_{j}" for j in range(len(doc_chunks))])
   chunk_sources.extend([doc["title"]] * len(doc_chunks))


print(f"Created {len(chunks)} chunks from {len(documents)} documents")

Шаг 5: Создание векторных представлений (эмбеддингов)

Используем предобученную модель sentence transformer от Hugging Face:

model_name = "sentence-transformers/all-MiniLM-L6-v2"
embedding_model = SentenceTransformer(model_name)


sample_text = "This is a sample text to test our embedding model."
sample_embedding = embedding_model.encode(sample_text)
print(f"Embedding dimension: {len(sample_embedding)}")

Шаг 6: Настройка ChromaDB

ChromaDB — это легковесная векторная база данных, идеально подходящая для нашей поисковой системы:

chroma_client = chromadb.Client()


embedding_function = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=model_name)


collection = chroma_client.create_collection(
   name="document_search",
   embedding_function=embedding_function
)


batch_size = 100
for i in range(0, len(chunks), batch_size):
   end_idx = min(i + batch_size, len(chunks))
  
   batch_ids = chunk_ids[i:end_idx]
   batch_chunks = chunks[i:end_idx]
   batch_sources = chunk_sources[i:end_idx]
  
   collection.add(
       ids=batch_ids,
       documents=batch_chunks,
       metadatas=[{"source": source} for source in batch_sources]
   )
  
   print(f"Added batch {i//batch_size + 1}/{(len(chunks)-1)//batch_size + 1} to the collection")


print(f"Total documents in collection: {collection.count()}")

Шаг 7: Поиск документов

Теперь переходим к самой интересной части — поиску по документам:

def search_documents(query, n_results=5):
   """
   Search for documents similar to the query.
  
   Args:
       query (str): The search query
       n_results (int): Number of results to return
  
   Returns:
       dict: Search results
   """
   start_time = time.time()
  
   results = collection.query(
       query_texts=[query],
       n_results=n_results
   )
  
   end_time = time.time()
   search_time = end_time - start_time
  
   print(f"Search completed in {search_time:.4f} seconds")
   return results


queries = [
   "What are the effects of climate change?",
   "History of artificial intelligence",
   "Space exploration missions"
]


for query in queries:
   print(f"\nQuery: {query}")
   results = search_documents(query)
  
   for i, (doc, metadata) in enumerate(zip(results['documents'][0], results['metadatas'][0])):
       print(f"\nResult {i+1} from {metadata['source']}:")
       print(f"{doc[:200]}...")

Шаг 8: Создание интерактивного интерфейса поиска

Создадим простую функцию для лучшего взаимодействия с пользователем:

def interactive_search():
   """
   Interactive search interface for the document search engine.
   """
   while True:
       query = input("\nEnter your search query (or 'quit' to exit): ")
      
       if query.lower() == 'quit':
           print("Exiting search interface...")
           break
          
       n_results = int(input("How many results would you like? "))
      
       results = search_documents(query, n_results)
      
       print(f"\nFound {len(results['documents'][0])} results for '{query}':")
      
       for i, (doc, metadata, distance) in enumerate(zip(
           results['documents'][0],
           results['metadatas'][0],
           results['distances'][0]
       )):
           relevance = 1 - distance  
           print(f"\n--- Result {i+1} ---")
           print(f"Source: {metadata['source']}")
           print(f"Relevance: {relevance:.2f}")
           print(f"Excerpt: {doc[:300]}...")  
           print("-" * 50)


interactive_search()

Шаг 9: Фильтрация результатов поиска по метаданным

Добавим возможность фильтрации результатов поиска по метаданным:

def filtered_search(query, filter_source=None, n_results=5):
   """
   Search with optional filtering by source.
  
   Args:
       query (str): The search query
       filter_source (str): Optional source to filter by
       n_results (int): Number of results to return
  
   Returns:
       dict: Search results
   """
   where_clause = {"source": filter_source} if filter_source else None
  
   results = collection.query(
       query_texts=[query],
       n_results=n_results,
       where=where_clause
   )
  
   return results


unique_sources = list(set(chunk_sources))
print(f"Available sources for filtering: {len(unique_sources)}")
print(unique_sources[:5])  


if len(unique_sources) > 0:
   filter_source = unique_sources[0]
   query = "main concepts and principles"
  
   print(f"\nFiltered search for '{query}' in source '{filter_source}':")
   results = filtered_search(query, filter_source=filter_source)
  
   for i, doc in enumerate(results['documents'][0]):
       print(f"\nResult {i+1}:")
       print(f"{doc[:200]}...")

Заключение

В этой статье мы рассмотрели, как построить семантическую поисковую систему для документов с использованием моделей встраивания от Hugging Face и ChromaDB. Система находит документы на основе их смысла, а не только по ключевым словам, путем преобразования текста в векторные представления.

Реализация включает обработку статей из Википедии, разделение их на фрагменты для более детального поиска, встраивание с помощью sentence transformers и хранение в векторной базе данных для эффективного извлечения. Конечный продукт предлагает интерактивный поиск, фильтрацию по метаданным и ранжирование по релевантности.

Такая система может быть применена для организации внутрикорпоративных знаний, анализа научных публикаций или любых других задач, требующих интеллектуального поиска по большому объему текстовых данных.