当前位置：首页 > article >正文

摸一下elasticsearch8的AI能力：语义搜索/vector向量搜索案例

article 2024/11/23 5:32:34

一、elasticsearch8.x+kibana docker-compose.yml一键安装

ES有RBAC安全验证,需要curl请求es注册用户并配置角色权限，比较麻烦，这里直接关闭,xpack.security.enabled=false 。ES版本8.15.3以上，否则执行搜索时可能会报错。

version: '3.3'

services:
  elasticsearch:
    image: docker.elastic.co/elasticsearch/elasticsearch:8.15.3
    container_name: elasticsearch
    environment:
      - discovery.type=single-node
      - bootstrap.memory_lock=true
      - "ES_JAVA_OPTS=-Xms512m -Xmx512m"
      - xpack.security.enabled=false
      - xpack.security.http.ssl.enabled=false
      - xpack.security.transport.ssl.enabled=false
      - ELASTIC_PASSWORD=123456
    ulimits:
      memlock:
        soft: -1
        hard: -1
    volumes:
      - esdata:/usr/share/elasticsearch/data
    ports:
      - "9200:9200"
      - "9300:9300"
    networks:
      - elastic

  kibana:
    image: docker.elastic.co/kibana/kibana:8.15.3
    container_name: kibana
    environment:
      - ELASTICSEARCH_HOSTS=http://elasticsearch:9200
    volumes:
      - /home/kibana.yml:/usr/share/kibana/config/kibana.yml
    ports:
      - "5601:5601"
    depends_on:
      - elasticsearch
    networks:
      - elastic

volumes:
  esdata:

networks:
  elastic:
    driver: bridge

/home/kibana.yml 就加了个汉化配置

server.host: "0.0.0.0"
server.shutdownTimeout: "5s"
elasticsearch.hosts: [ "http://elasticsearch:9200" ]
i18n.locale: "zh-CN"

二、语料向量化模型选择
modelscope下载一个适合的预训练模型用于将商品信息,文章摘要等资源类信息向量化后存入ES。模型都从modelcope找的，hugginface也能找到差不多的，但需要代理且速度很慢。
先试了bert-base-chinese和bert-base-uncased,但效果很不好，他们不会认为高兴和愉快是近义词，而会认为高兴和高处，很好和很不好是近义词，也就是有字词重叠的词语间向量化后的欧式距离更小。
实现语义搜索应使用sentence BERT类模型，这里找到了nlp_gte_sentence-embedding_chinese-large，效果就很不错。
三、模型及python依赖下载

pip3 install modelscope torch transformers  elasticsearch
modelscope download --model iic/nlp_gte_sentence-embedding_chinese-large

如果有如下SSL问题，可以强行改下python requests这个标准库的sessions.py 779行左右的verify=False
在这里插入图片描述

四、语义搜索案例
大致逻辑：

transformers库加载模型及分词器
将商品信息,文章摘要等资源类信息向量化后存入ES
搜索input案例也使用同样方式向量化调用ES script query API

具体代码如下，去除空行及注释大概就100行

from transformers import BertTokenizer, BertModel
import torch
import time
from elasticsearch import Elasticsearch

# 初始化模型和分词器
MODEL_PATH = 'C:\\Users\\Administrator\\.cache\\modelscope\\hub\\iic\\nlp_gte_sentence-embedding_chinese-large'
#MODEL_PATH = 'C:\\Users\\Administrator\\.cache\\modelscope\\hub\\tiansz\\bert-base-chinese'
#MODEL_PATH = 'C:\\Users\\Administrator\\.cache\\modelscope\\hub\\AI-ModelScope\\bert-base-uncased'

tokenizer = BertTokenizer.from_pretrained(MODEL_PATH)
model = BertModel.from_pretrained(MODEL_PATH)
es = Elasticsearch([{'scheme':'http','host':'192.168.72.128','port':9200}])

def embed_texts(texts, batch_size=32):
    """
    为一组文本生成BERT嵌入向量。

    :param texts: 需要生成嵌入向量的文本列表
    :param batch_size: 每次处理的文本数量，默认为32
    :return: 嵌入向量列表
    """
    all_embeddings = []
    for i in range(0, len(texts), batch_size):
        batch_texts = texts[i:i+batch_size]
        inputs = tokenizer(batch_texts, return_tensors='pt', padding=True, max_length=512, truncation=True)
        with torch.no_grad():
            outputs = model(**inputs)
            embeddings = outputs.last_hidden_state.mean(dim=1).detach().numpy()
            all_embeddings.extend(embeddings)
    return all_embeddings

def index_document(doc_id, text):
    """
    将文档及其嵌入向量索引到Elasticsearch中。

    :param doc_id: 文档的唯一标识符
    :param text: 文档的文本内容
    """
    try:
        embedding = embed_texts([text])[0]
        doc = {'id': doc_id, 'text': text, 'embedding': embedding.tolist()}
        res = es.index(index='documents', id=doc_id, body=doc)
        if res['result'] != 'created':
            print(f"Failed to index document {doc_id}")
    except Exception as e:
        print(f"Error indexing document {doc_id}: {e}")

def search_similar_documents(query, top_k=10):
    """
    根据查询文本搜索最相似的文档。

    :param query: 查询文本
    :param top_k: 返回的最相似文档数量，默认为3
    :return: 最相似的文档列表
    """
    query_embedding = embed_texts([query])[0]

    script_query = {
        "script_score": {
            "query": {"match_all": {}},
            "script": {
                "source": "cosineSimilarity(params.query_vector, 'embedding') + 1.0",
                "params": {"query_vector": query_embedding.tolist()}
            }
        }
    }
    body = {
        "size": top_k,
        "query": script_query,
        "_source": {"includes": ["text","_score"]}
    }

    response = es.search(index='documents', body=body)
    return response['hits']['hits']

# 示例文档
documents = [
    {"id": 1, "text": "今天天气真好"},
    {"id": 2, "text": "今儿天气很好"},
    {"id": 3, "text": "这里气候很好"},
    {"id": 4, "text": "今天天气很差"},
    {"id": 5, "text": "今天天气很不好"},
    {"id": 6, "text": "今天天气真差"}
]


# 将示例文档索引到Elasticsearch中
for doc in documents:
    index_document(doc['id'], doc['text'])

# 示例查询
#ES是AP模型，需要等下再查，否则查不到数据，着急的话可以去kibana的dev console跑一下`POST documents/_refresh`
time.sleep(3)
results = search_similar_documents('今天天气差极了')
print(results)

效果符合预期，今天天气差极了相关度高的几个sentence的_score最高的排在了最前面，今儿天气很好虽然描述天气好，但至少与天气相关，而这里气候很好描述的某地的气候，确实为最不相关，这种效果是ES传统的term,fuzzy,prefix,phase_match,regex等搜索方式达不到的。
在这里插入图片描述