Используйте LLM для поиска документов в Elasticsearch

Семантический поиск был модной темой в недавней шумихе вокруг больших языковых моделей (LLM). Основная причина заключается в том, что в традиционном текстовом поиске используется лексический подход, имеющий известные ограничения. Он ищет буквальные совпадения (или их варианты, такие как версия с основой) слов запроса, введенных пользователем. Этот подход упускает контекст и не понимает, что на самом деле означает весь запрос.

Например, когда пользователь выполняет поиск по слову страховка, решение для лексического поиска не сможет найти документы, содержащие слово Медикэйд, но не содержащие слова страхование в явном виде. Быстрое и простое решение этой проблемы — воспользоваться преимуществами мощных LLM с открытым исходным кодом, доступных на хабе моделей Huggingface.

В этом посте и сопутствующем блокноте я показываю

Вот шаги высокого уровня:

  • Выберите стратегию токенизации, например, подходит ли вам встраивание на уровне слова, предложения или документа.
  • Создавайте вложения на уровне документа, учитывая, что многие базовые модели подходят для семантического поиска и представления предложений в целом.
  • Запустите кластер Elasticsearch с одним узлом локально, следуя документации по настройке докера.
  • Индексируйте документы и их вложения.
  • Наконец, сгенерируйте векторы запросов, используя один и тот же кодировщик во время поиска, чтобы получить векторы, наиболее похожие на вектор запроса.

Для целей этого сообщения я использую набор данных группы новостей в качестве корпуса, маленькую модель T5 для получения вложений на уровне документа и локально запускаемый Кластер эластичного поиска. Полный блокнот находится на моем GitHub.

Обратите внимание, что вы всегда можете выбрать лучшую модель (например, FLAN-T5) и настроить выбранную модель на свой собственный корпус, чтобы получить лучшие результаты. Вам также следует обратить внимание на показатели ранжированного поиска, такие как nDCG, чтобы оценить, насколько хорошо работает ваша система семантического поиска.

1. Настройте кластер Elasticsearch с одним узлом в локальной среде.

Вы можете пропустить этот раздел, если он у вас уже есть.

  1. извлеките образ Elasticsearch Docker docker pull docker.elastic.co/elasticsearch/elasticsearch:8.7.0
  2. Создайте новую сеть докеров для Elasticsearch и Kibana docker network create elastic
  3. Запустите Elasticsearch в Docker docker run --name es01 --net elastic -p 9200:9200 -it docker.elastic.co/elasticsearch/elasticsearch:8.7.0
  4. Скопируйте сертификат безопасности http_ca.crt из контейнера Docker на локальный компьютер docker cp es01:/usr/share/elasticsearch/config/certs/http_ca.crt .
  5. [необязательно] настройте Kibana: docker run \ --name kibana \ --net elastic \ -p 5601:5601 \ docker.elastic.co/kibana/kibana:8.2.2 обратите внимание, что вам понадобится токен регистрации из шага 3
  6. Убедитесь, что у вас есть доступ к кластеру из записной книжки. Подробнее о Elasticsearch python cli.

2. Создайте индекс Elasticsearch

from elasticsearch import Elasticsearch

ELASTIC_PASSWORD = "PASSWORD"
ES_HOST = "https://localhost:9200/"
index_name = "semantic-search"

# Create the client instance
client = Elasticsearch(
    hosts=ES_HOST,
    ca_certs='./http_ca.crt',
    basic_auth=("elastic", ELASTIC_PASSWORD)
)

# get cluster information
client.info()

# define index config
config = {
    "mappings": {
        "properties": {
            "text": {"type": "text"},
            "embeddings": {
                    "type": "dense_vector",
                    "dims": 512,
                    "index": True
                }
            }
    },
    "settings": {
        "number_of_shards": 2,
        "number_of_replicas": 1
    }
}

# create an index in elasticsearch
try:
    client.indices.create(
        index=index_name,
        settings=config["settings"],
        mappings=config["mappings"],
    )
except:
    print(f"Index already exists: {client.indices.exists(index=[index_name])}")

3. Получите векторы встраивания

import torch
from datasets import load_dataset
from transformers import T5Model, T5Tokenizer

# load the dataset
dataset = load_dataset('newsgroup', '18828_alt.atheism')
# check an example of data
dataset['train'][0]['text']

# load the pre-trained model
tokenizer = T5Tokenizer.from_pretrained("t5-small")
model = T5Model.from_pretrained("t5-small")

# function to get embeddings
def get_embeddings(input_text, model=model, tokenizer=tokenizer, max_length=512):
        
    inputs = tokenizer.encode_plus(input_text, 
                                         max_length=max_length,
                                         pad_to_max_length=True,
                                         return_tensors="pt")
    
    outputs = model(input_ids=inputs['input_ids'], decoder_input_ids=inputs['input_ids'])
    
    last_hidden_states = torch.mean(outputs[0], dim=1)

    return last_hidden_states.tolist()

# index documents and their embedding in Elasticsearch
for i in range(small_dataset.num_rows):
    doc = {"text": small_dataset['text'][i],
           "embeddings": get_embeddings(small_dataset['text'][i])[0]
    }
    
    client.index(index= index_name, document=doc)

# check the number of saved documents
result = client.count(index=index_name)
print(result.body['count'])

4. Ищите!

После завершения индексации мы можем искать наши данные. Elasticsearch использует косинусное сходство, но также предоставляет оболочку Python для выполнения поиска KNN. У вас также есть возможность использовать пользовательскую функцию сходства.

Я предоставил фрагмент кода для KNN с k = 5. Помните, что вы также должны предоставить вложения для вашего термина запроса.

query_embedding = get_embeddings(dataset['train']['text'][20])[0]
query_dict = {
    "field": "embeddings",
    "query_vector": query_embedding,
    "k": 5,
    "num_candidates": 5
}
res = client.knn_search(index=index_name, knn=query_dict, source=["text"])

Дополнительные материалы на PlainEnglish.io.

Подпишитесь на нашу бесплатную еженедельную рассылку новостей. Подпишитесь на нас в Twitter, LinkedIn, YouTube и Discord .