Clever Geek Handbook

Comment rechercher des marais de fichiers dans 104 lignes de code en python

Poursuivant le sujet des courts scripts utiles, je voudrais faire connaître aux lecteurs la possibilité de construire une recherche par le contenu des fichiers et des images en 104 lignes. Ce ne sera certainement pas une solution époustouflante - mais cela fonctionnera pour des besoins simples. De plus, l'article n'inventera rien - tous les packages sont open source.



Et oui - les lignes vides dans le code sont également comptées. Une petite démonstration de travail est donnée à la fin de l'article.



Nous avons besoin de python3 , téléchargé par Tesseract 5, et du modèle distiluse-base-multilingual-cased du package Sentence-Transformers . Ceux qui comprennent déjà ce qui va se passer ensuite ne seront pas intéressants.



En attendant, tout ce dont nous avons besoin ressemblera à:



18 premières lignes 
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')




Il sera nécessaire, comme vous pouvez le voir, décemment, et tout semble prêt, mais vous ne pouvez pas vous passer d'un fichier. En particulier, textract (pas d'Amazon, qui est payé), ne fonctionne pas bien avec les fichiers PDF russes, car vous pouvez utiliser pdfplumber . De plus, diviser le texte en phrases est une tâche difficile et, dans ce cas, razdel fait un excellent travail avec la langue russe .



Ceux qui n'ont pas entendu parler de scikit-learn - j'envie qu'en bref, l'algorithme NearestNeighbors qu'il contient se souvienne des vecteurs et donne les plus proches. Au lieu de scikit-learn, vous pouvez utiliser faiss ou ennuyer, ou même elasticsearch par exemple .



L'essentiel est de transformer le texte de (n'importe quel) fichier en vecteur, ce qu'ils font:



36 lignes de code suivantes 
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
        )]
    }




Eh bien, cela reste une question de technique - parcourir tous les fichiers, extraire les vecteurs et trouver le plus proche de la requête par la distance cosinus.



Code restant 
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(' : ')




Vous pouvez exécuter tout le code comme ceci: 



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


C'est comme ça avec le code.



Et voici la démo promise.



J'ai pris deux nouvelles de "Lenta.ru", et mis l'une dans un fichier gif à travers la peinture notoire, et l'autre juste dans un fichier texte.



Fichier First.gif




Deuxième fichier .txt
, . .



, - . , , , . . , .



, , , . . .



, - - .



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



Et voici une animation gif de son fonctionnement. Avec le GPU, bien sûr, tout fonctionne plus joyeusement.



Démonstration, mieux cliquer sur l'image






Merci d'avoir lu! J'espère toujours que cette méthode sera utile à quelqu'un.


More articles:


All Articles