Cómo buscar pantanos de archivos en 104 líneas de código en Python

Continuando con el tema de los guiones breves y útiles, me gustaría familiarizar a los lectores con la posibilidad de construir una búsqueda por el contenido de archivos e imágenes en 104 líneas. Ciertamente, esta no será una solución alucinante, pero funcionará para necesidades simples. Además, el artículo no inventará nada: todos los paquetes son de código abierto.



Y sí, también se cuentan las líneas en blanco en el código. Se da una pequeña demostración del trabajo al final del artículo.



Necesitamos python3 , descargado por Tesseract 5, y el modelo distiluse-base-multilingual-cased del paquete Sentence-Transformers . Quien ya entienda lo que sucederá a continuación, no será interesante.



Mientras tanto, todo lo que necesitamos se verá así:



Primeras 18 líneas
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')





Será necesario, como puede ver, decentemente, y todo parece estar listo, pero no puede prescindir de un archivo. En particular, textract (no de Amazon, que es de pago), de alguna manera no funciona bien con archivos PDF rusos, ya que puede usar pdfplumber . Además, dividir el texto en oraciones es una tarea difícil y, en este caso, razdel hace un excelente trabajo con el idioma ruso .



Aquellos que no han oído hablar de scikit-learn - Envidio que, en resumen, el algoritmo NehestNeighbors en él recuerde los vectores y proporcione los más cercanos. En lugar de scikit-learn, puede utilizar Faiss o molestar, o incluso Elasticsearch por ejemplo .



Lo principal es convertir el texto de (cualquier) archivo en un vector, que es lo que hacen:



siguientes 36 líneas de código
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
        )]
    }





Bueno, entonces sigue siendo una cuestión de técnica: revisar todos los archivos, extraer los vectores y encontrar el más cercano a la consulta por la distancia del coseno.



Código restante
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(' : ')





Puedes ejecutar todo el código así:



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


Así es con el código.



Y aquí está la demostración prometida.



Tomé dos noticias de "Lenta.ru", y puse una en un archivo gif a través de la notoria pintura, y la otra solo en un archivo de texto.



Archivo First.gif




Segundo archivo .txt
, . .



, - . , , , . . , .



, , , . . .



, - - .



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



Y aquí hay una animación gif de cómo funciona. Con la GPU, por supuesto, todo funciona más alegre.



Demostración, mejor haga clic en la imagen






¡Gracias por leer! Todavía espero que este método sea útil para alguien.



All Articles