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

    Продолжая тематику коротких полезных скриптов, хотелось бы познакомить читателей с возможностью построения поиска по контенту файлов и изображений в 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 конечно всё работает бодрее.

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

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

    Similar posts

    AdBlock has stolen the banner, but banners are not teeth — they will be back

    More
    Ads

    Comments 5

      +2

      в функции 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':, и не повторять кучу раз код. Да, три лишние строчки добавится, но код станет читабельнее.

        +1

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


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

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

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


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


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

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

          Only users with full accounts can post comments. Log in, please.