Spaces:
Sleeping
Sleeping
File size: 4,874 Bytes
57cf043 |
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 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 |
import json
import logging
from pathlib import Path
import time
from elasticsearch import Elasticsearch
from tqdm import tqdm
def create_index_elastic_rocks_nn(
path: str,
logger: logging.Logger | None = None,
):
if logger is None:
logger = logging.getLogger(__name__)
# Подключение к Elasticsearch
es = Elasticsearch(hosts='localhost:9200')
INDEX_NAME = 'rocks_nn_search_elastic'
# Удаление старого индекса, если он существует
if es.indices.exists(index=INDEX_NAME):
es.indices.delete(index=INDEX_NAME)
mapping = {
"settings": {
"analysis": {
"filter": {
"custom_stopwords": {
"type": "stop",
"stopwords": [
"ООО",
"ОАО",
"НН",
"нн",
"Перечень",
"перечень",
"дивизиона",
"дивизион",
],
}
},
"analyzer": {
"custom_analyzer": {
"type": "custom",
"tokenizer": "standard",
"filter": [
"lowercase",
"custom_stopwords",
],
}
},
}
},
"mappings": {
"properties": {
"division_name": {
"type": "text",
"analyzer": "custom_analyzer",
"search_analyzer": "custom_analyzer",
},
"division_name_2": {
"type": "text",
"analyzer": "custom_analyzer",
"search_analyzer": "custom_analyzer",
},
"company_name": {
"type": "text",
"analyzer": "custom_analyzer",
"search_analyzer": "custom_analyzer",
},
}
},
}
# Создание индекса с указанным маппингом
es.indices.create(index=INDEX_NAME, body=mapping)
for ind, path in tqdm(enumerate(Path(path).iterdir())):
# Открываем файл и читаем его содержимое
with open(path, 'r', encoding='utf-8') as file:
data = json.load(file)
# Индексирование документа в Elasticsearch
es.index(index=INDEX_NAME, id=ind + 1, body=data)
# Подсчет количества документов в индексе
count_response = es.count(index=INDEX_NAME)
logger.info(
f"{ind}, Total documents in '{INDEX_NAME}': {count_response['count']}"
)
time.sleep(1.0)
if es.indices.exists(index=INDEX_NAME):
logger.info(f"Index '{INDEX_NAME}' exists.")
# Подсчет количества документов в индексе
count_response = es.count(index=INDEX_NAME)
logger.info(f"Total documents in '{INDEX_NAME}': {count_response['count']}")
query = "Какие РОКС НН входят в состав Норильского дивизиона?"
query_ = {
"query": {
"function_score": {
"query": {
"multi_match": {
"query": f"{query}",
"fields": ["division_name", "division_name_2", "company_name"],
"fuzziness": "AUTO",
"analyzer": "custom_analyzer",
}
},
"functions": [
{
"filter": {
"term": {"_id": "3"} # ID документа, который нужно понизить
},
"weight": 0.5, # Устанавливает очень низкий вес для этого документа
}
],
"boost_mode": "multiply", # Сочетание _score и весов
}
}
}
# Выполнение поиска в Elasticsearch
response = es.search(index=INDEX_NAME, body=query_, size=1)
logger.info(f"Number of hits: {response['hits']['total']['value']}")
# Вывод результата поиска
for hit in response['hits']['hits']:
logger.info(hit['_source'])
if __name__ == '__main__':
path = '/mnt/ntr_work/project/nmd800/data/rocks_nn_card'
create_index_elastic_rocks_nn(path)
|