Cara mencari file swamps dalam 104 baris kode di python

Melanjutkan topik skrip singkat yang berguna, saya ingin memperkenalkan pembaca dengan kemungkinan membangun pencarian berdasarkan konten file dan gambar dalam 104 baris. Ini tentunya bukan solusi yang membingungkan - tetapi akan berhasil untuk kebutuhan sederhana. Selain itu, artikel tidak akan menemukan apa pun - semua paket bersifat open source.



Dan ya - baris kosong dalam kode juga dihitung. Demonstrasi kecil tentang kerja diberikan di akhir artikel.



Kami membutuhkan python3 , diunduh oleh Tesseract 5, dan model berbasis multibahasa distiluse-base- dari paket Sentence-Transformers . Mereka yang sudah mengerti apa yang akan terjadi selanjutnya tidak akan menarik.



Sementara itu, semua yang kami butuhkan akan terlihat seperti:



18 baris pertama
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')





Ini akan dibutuhkan, seperti yang Anda lihat, dengan sopan, dan semuanya tampaknya siap, tetapi Anda tidak dapat melakukannya tanpa file. Secara khusus, textract (bukan dari Amazon, yang berbayar), entah bagaimana tidak berfungsi dengan baik dengan pdf Rusia, karena Anda dapat menggunakan pdfplumber . Lebih lanjut, membagi teks menjadi beberapa kalimat adalah tugas yang sulit, dan dalam hal ini razdel melakukan pekerjaan yang sangat baik dengan bahasa Rusia .



Mereka yang belum pernah mendengar tentang scikit-learn - Saya iri karena singkatnya, algoritme NearestNeighbours di dalamnya mengingat vektor dan memberikan yang terdekat. Alih-alih scikit-learn, Anda dapat menggunakan faiss atau mengganggu atau bahkan elasticsearch misalnya .



Hal utama adalah benar-benar mengubah teks dari file (apa saja) menjadi vektor, itulah yang mereka lakukan:



36 baris kode berikutnya
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
        )]
    }





Jadi, tinggal masalah teknik - untuk memeriksa semua file, mengekstrak vektor dan menemukan yang paling dekat dengan kueri dengan jarak cosinus.



Kode yang tersisa
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(' : ')





Anda dapat menjalankan semua kode seperti ini:



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


Begitulah halnya dengan kode.



Dan inilah demo yang dijanjikan.



Saya mengambil dua berita dari "Lenta.ru", dan memasukkan satu ke dalam file gif melalui cat terkenal, dan yang lainnya hanya ke dalam file teks.



File first.gif




File .txt kedua
, . .



, - . , , , . . , .



, , , . . .



, - - .



, №71 , , , . 10 , . — .



Dan berikut ini adalah animasi gif cara kerjanya. Dengan GPU, tentunya semuanya bekerja lebih ceria.



Demonstrasi, lebih baik klik gambarnya






Terima kasih sudah membaca! Saya masih berharap cara ini bermanfaat bagi seseorang.



All Articles