Используйте LLM для поиска документов в Elasticsearch
Семантический поиск был модной темой в недавней шумихе вокруг больших языковых моделей (LLM). Основная причина заключается в том, что в традиционном текстовом поиске используется лексический подход, имеющий известные ограничения. Он ищет буквальные совпадения (или их варианты, такие как версия с основой) слов запроса, введенных пользователем. Этот подход упускает контекст и не понимает, что на самом деле означает весь запрос.
Например, когда пользователь выполняет поиск по слову страховка, решение для лексического поиска не сможет найти документы, содержащие слово Медикэйд, но не содержащие слова страхование в явном виде. Быстрое и простое решение этой проблемы — воспользоваться преимуществами мощных LLM с открытым исходным кодом, доступных на хабе моделей Huggingface.
В этом посте и сопутствующем блокноте я показываю
Вот шаги высокого уровня:
- Выберите стратегию токенизации, например, подходит ли вам встраивание на уровне слова, предложения или документа.
- Создавайте вложения на уровне документа, учитывая, что многие базовые модели подходят для семантического поиска и представления предложений в целом.
- Запустите кластер Elasticsearch с одним узлом локально, следуя документации по настройке докера.
- Индексируйте документы и их вложения.
- Наконец, сгенерируйте векторы запросов, используя один и тот же кодировщик во время поиска, чтобы получить векторы, наиболее похожие на вектор запроса.
Для целей этого сообщения я использую набор данных группы новостей в качестве корпуса, маленькую модель T5 для получения вложений на уровне документа и локально запускаемый Кластер эластичного поиска. Полный блокнот находится на моем GitHub.
Обратите внимание, что вы всегда можете выбрать лучшую модель (например, FLAN-T5) и настроить выбранную модель на свой собственный корпус, чтобы получить лучшие результаты. Вам также следует обратить внимание на показатели ранжированного поиска, такие как nDCG, чтобы оценить, насколько хорошо работает ваша система семантического поиска.
1. Настройте кластер Elasticsearch с одним узлом в локальной среде.
Вы можете пропустить этот раздел, если он у вас уже есть.
- извлеките образ Elasticsearch Docker
docker pull docker.elastic.co/elasticsearch/elasticsearch:8.7.0
- Создайте новую сеть докеров для Elasticsearch и Kibana
docker network create elastic
- Запустите Elasticsearch в Docker
docker run --name es01 --net elastic -p 9200:9200 -it docker.elastic.co/elasticsearch/elasticsearch:8.7.0
- Скопируйте сертификат безопасности http_ca.crt из контейнера Docker на локальный компьютер
docker cp es01:/usr/share/elasticsearch/config/certs/http_ca.crt .
- [необязательно] настройте Kibana:
docker run \ --name kibana \ --net elastic \ -p 5601:5601 \ docker.elastic.co/kibana/kibana:8.2.2
обратите внимание, что вам понадобится токен регистрации из шага 3 - Убедитесь, что у вас есть доступ к кластеру из записной книжки. Подробнее о 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 .