Как сделать поиск по файлоболотам в 104 строки кода на python +15



Продолжая тематику коротких полезных скриптов, хотелось бы познакомить читателей с возможностью построения поиска по контенту файлов и изображений в 104 строки. Это конечно не будет умопомрачительным по качеству решением — но вполне годным для простых нужд. Также в статье не будет ничего изобретаться — все пакеты open source.

И да — пустые строки в коде тоже считаются. Небольшая демонстрация работы приведена в конце статьи.

Нам понадобится python3, скачанный Tesseract пятой версии, и модель distiluse-base-multilingual-cased из пакета Sentence-Transformers. Кому уже понятно что дальше будет происходить — интересно не будет.

А тем временем, всё что нам понадобится, будет выглядеть как:

Первые 18 строк
import numpy as np
import os, sys, glob

os.environ['PATH'] += os.pathsep + os.path.join(os.getcwd(), 'Tesseract-OCR')
extensions = [
    '.xlsx', '.docx', '.pptx',
    '.pdf', '.txt', '.md', '.htm', 'html',
    '.jpg', '.jpeg', '.png', '.gif'
]

import warnings; warnings.filterwarnings('ignore')
import torch, textract, pdfplumber
from cleantext import clean
from razdel import sentenize
from sklearn.neighbors import NearestNeighbors
from sentence_transformers import SentenceTransformer
embedder = SentenceTransformer('./distillUSE')



Понадобится как видим, прилично, и вроде всё готовое, но и без напильника не обойтись. В частности, textract (не от Amazon который платный), как-то плохо работает с русскими pdf, как выход использовать можно pdfplumber. Далее, разбиение текста на предложения — сложная задача, и с русским языком в её случае отлично справляется razdel.

Кто не слышал про scikit-learnтому я завидую вкратце, алгоритм NearestNeighbors в нём запоминает вектора и выдает ближайшие. Вместо scikit-learn можно использовать faiss или annoy или например даже elasticsearch.

Главное на самом деле превратить текст (любого) файла в вектор, что и делают:

следующие 36 строк кода
def processor(path, embedder):
    try:
        if path.lower().endswith('.pdf'):
            with pdfplumber.open(path) as pdf:
                if len(pdf.pages):
                    text = ' '.join([
                        page.extract_text() or '' for page in pdf.pages if page
                    ])
        elif path.lower().endswith('.md') or path.lower().endswith('.txt'):
            with open(path, 'r', encoding='UTF-8') as fd:
                text = fd.read()
        else:
            text = textract.process(path, language='rus+eng').decode('UTF-8')
        if path.lower()[-4:] in ['.jpg', 'jpeg', '.gif', '.png']:
            text = clean(
                text,
                fix_unicode=False, lang='ru', to_ascii=False, lower=False,
                no_line_breaks=True
            )
        else:
            text = clean(
                text,
                lang='ru', to_ascii=False, lower=False, no_line_breaks=True
            )
        sentences = list(map(lambda substring: substring.text, sentenize(text)))
    except Exception as exception:
        return None
    if not len(sentences):
        return None
    return {
        'filepath': [path] * len(sentences),
        'sentences': sentences,
        'vectors': [vector.astype(float).tolist() for vector in embedder.encode(
            sentences
        )]
    }



Ну а далее остаётся дело техники — обойти все файлы, извлечь вектора и найти ближайший к запросу по косинусному расстоянию.

Оставшийся код
def indexer(files, embedder):
    for file in files:
        processed = processor(file, embedder)
        if processed is not None:
            yield processed

def counter(path):
    if not os.path.exists(path):
        return None
    for file in glob.iglob(path + '/**', recursive=True):
        extension = os.path.splitext(file)[1].lower()
        if extension in extensions:
            yield file

def search(engine, text, sentences, files):
    indices = engine.kneighbors(
        embedder.encode([text])[0].astype(float).reshape(1, -1),
        return_distance=True
    )

    distance = indices[0][0][0]
    position = indices[1][0][0]

    print(
        'Релевантность "%.3f' % (1 - distance / 2),
        'Фраза: "%s", файл "%s"' % (sentences[position], files[position])
    )

print('Поиск файлов "%s"' % sys.argv[1])
paths = list(counter(sys.argv[1]))

print('Индексация "%s"' % sys.argv[1])
db = list(indexer(paths, embedder))

sentences, files, vectors = [], [], []
for item in db:
    sentences += item['sentences']
    files += item['filepath']
    vectors += item['vectors']

engine = NearestNeighbors(n_neighbors=1, metric='cosine').fit(
    np.array(vectors).reshape(len(vectors), -1)
)

query = input('Что искать: ')
while query:
    search(engine, query, sentences, files)
    query = input('Что искать: ')



Запускать весь код можно так:

python3 app.py /path/to/your/files/

Вот как бы и всё с кодом.

А вот обещанная демонстрация.

Взял две новости с «лента.ру», и положил одну в gif-файл через небезызвестный paint, а вторую просто в текстовый.

Первый файл.gif


Второй файл.txt
Специалисты МЧС дали россиянам рекомендации, как снизить риск попадания в ДТП. Об этом сообщает РИА Новости.

Как отметили в ведомстве, большинство аварий происходит из-за беспечности и неправильного поведения на дороге. Поэтому, двигаясь по трассе, необходимо соблюдать предельную внимательность и использовать периферийное зрение, то есть смотреть в одну точку. Это позволяет видеть больше объектов по бокам. Важно внимательно следить за едущими рядом и сзади автомобилями, их скоростью.

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

В МЧС добавили, что чаще всего аварии происходят из-за управления автомобилем в нетрезвом виде и из-за превышения скорости.

Ранее доктор Александр Мясников, главврач ГКБ №71 имени Жадкевича в Москве, назвал людей, которые пользуются мобильными телефонами за рулем, преступниками и сравнил их с пьяными водителями. Он предложил штрафовать таких водителей на 10 тысяч рублей после того, как пользование мобильным телефоном зафиксировали камеры видеонаблюдения. За повторное нарушение — лишение прав.

А вот gif-анимация, как это работает. С GPU конечно всё работает бодрее.

Демонстрация, лучше кликнуть на картинку

Спасибо за ознакомление! Надеюсь всё же что этот метод кому-нибудь да будет полезен.




Комментарии (5):

  1. azxc
    /#21914214 / +1

    в функции processor:


    try:
    ...
    except Exception as exception:

    не позволил вам найти опечатку в строчке


    with pdfplumber.open(path) as pfd:

    тут явная опечатка pfd должно быть pdf. ну и, кмк if len(pdf.pages): можно заменить было бы на if pdf.pages:.


    А еще там же, можно было бы написать


    from pathlib import Path
    
    ext = Path(path).suffix.lower()

    и дальше сравнивать if ext == '.pdf':, и не повторять кучу раз код. Да, три лишние строчки добавится, но код станет читабельнее.

    • S_A
      /#21914218

      Огонь, спасибо за ценное замечание и внимательность! Опечатку поправлю как доберусь до десктопа.


      Изначально при отладке в except стоял print(exception) и он был выпилен. with появился чуть позже. Спасибо еще раз.

    • S_A
      /#21914244

      Исправил опечатку и проверил что работает с PDF нормально. Про пути оставил из соображений сохранения оригинального кода. Но впредь обещаю себе больше не заниматься оптимизациями количества строк.

      • azxc
        /#21914260 / +1

        кмк, лучше переписать список ожидаемых exceptions, чем Exception. Но это дело вкуса.


        я в последнее время очень полюбил модуль pathlib (раньше использовал py.path), и всем советую его. но если не хочется, всегда можно сделать os.path.splitext(filename)[1]. собственно идея комментария была в том, что бы не делать path.lower() несколько раз в коде + сделать аккуратнее логику с .jpg/jpeg.


        а вообще спасибо за статью — мне было интересно почитать и кое-что новое узнал. когда-нибудь пригодится :)

  2. caxap1
    /#21916274

    Очень интересно, так как ежедневно работаю с большим количеством pdf файлов, в том числе с поиском по тексту файлов с помощью гугл диска.
    На самом деле, как оказалось, на рынке практически отсутствуют вменяемые решения каталогизации и поиска по содержимому pdf файлов.
    Буду признателен если кто-то может порекомендовать коробочные решения по хранению и интеллектуальному поиску по тексту большой массы pdf файлов.