Changes

7 changed files (+158/-171)

• .gitmodules (deleted)
  • @@ -1,3 +0,0 @@
    - [submodule "ImageBind"] - path = ImageBind - url = https://github.com/facebookresearch/ImageBind.git
• README.md
  • @@ -1,6 +1,6 @@
    # search - Semantic file search using ImageBind and sqlite-vec + Semantic file search using FastEmbed and sqlite-vec ## Installation
  • @@ -8,13 +8,13 @@ ```
    pip install -r requirements.txt ``` - If you don't need PyTorch with GPU support, first run `pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu`. During runtime, if you get the error `ModuleNotFoundError: No module named 'torchvision.transforms.functional_tensor'`, change `functional_tensor` to `functional`. + For GPU support, install `fastembed-gpu`. ## Usage - This program uses a client-server architecture to watch directories with inotify and keep the model loaded in memory. It takes around 15 seconds to load the model so it wouldn't be great if every query had to wait on that. + This program uses a client-server architecture to watch directories with inotify and keep the model in memory so the client doesn't have to wait several seconds to load the model. It uses file inodes and modification times to avoid unnecessary re-indexing. - Run `python server.py DIRS_TO_INDEX` to start the server. The server only indexes images, audio, and videos since text files tend to irrelevantly pollute the search results. + Run `python server.py DIRS_TO_INDEX` to start the server. Make sure that none of the directories contain other directories on that list or weird stuff will happen. The server currently only indexes images although more modalities may be supported in the future. There are probably some weird race condition bugs if you modify a lot of files at the same time. Then run `python client.py SEARCH_TEXT NUM_RESULTS` to get a list of the most similar files. You can pass this list to an image viewer such as Gwenview to view image results. Note that Gwenview doesn't preserve the order of the images. Alternatively, add a third parameter to `python client.py` and it will symlink `res0`, `res1`, and so on to the files on the list. This can be used in conjunction with the Dolphin file manager's integrated terminal to get thumbnails of the search results.
• client.py
  • @@ -1,23 +1,7 @@
    - from http.client import HTTPConnection import os - import socket import sys import xmlrpc.client - - - class UnixStreamHTTPConnection(HTTPConnection): - def connect(self): - self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) - self.sock.connect(self.host) - - - class UnixStreamTransport(xmlrpc.client.Transport, object): - def __init__(self, socket_path): - self.socket_path = socket_path - super().__init__() - - def make_connection(self, host): - return UnixStreamHTTPConnection(self.socket_path) + from unixsocket import UnixStreamTransport proxy = xmlrpc.client.ServerProxy(
• model.py
  • @@ -1,38 +1,16 @@
    - import torch - from imagebind import data - from imagebind.models import imagebind_model - from imagebind.models.imagebind_model import ModalityType + from fastembed import ImageEmbedding, TextEmbedding + import pillow_avif - device = "cuda:0" if torch.cuda.is_available() else "cpu" print("Loading model") - model = imagebind_model.imagebind_huge(pretrained=True) - model.eval() - model.to(device) - - def embed(mode, data): - with torch.no_grad(): - return model({mode: data})[mode][0] + text_model = TextEmbedding(model_name="jinaai/jina-clip-v1") + image_model = ImageEmbedding(model_name="jinaai/jina-clip-v1") def embed_text(text): - return embed(ModalityType.TEXT, data.load_and_transform_text([text], device)) - - - def embed_audio(audio_path): - return embed( - ModalityType.AUDIO, data.load_and_transform_audio_data([audio_path], device) - ) + return next(text_model.embed(text)) def embed_image(image_path): - return embed( - ModalityType.VISION, data.load_and_transform_vision_data([image_path], device) - ) - - - def embed_video(video_path): - return embed( - ModalityType.VISION, data.load_and_transform_video_data([video_path], device) - ) + return next(image_model.embed(image_path))
• requirements.txt
  • @@ -1,4 +1,30 @@
    - file:ImageBind + certifi==2024.12.14 + charset-normalizer==3.4.1 + coloredlogs==15.0.1 + fastembed==0.5.0 + filelock==3.16.1 + flatbuffers==24.12.23 + fsspec==2024.12.0 + huggingface-hub==0.27.1 + humanfriendly==10.0 + idna==3.10 + loguru==0.7.3 + mmh3==4.1.0 + mpmath==1.3.0 + numpy==2.2.1 + onnx==1.17.0 + onnxruntime==1.20.1 + packaging==24.2 + pillow==10.4.0 pillow-avif-plugin==1.4.6 - watchdog==4.0.1 - sqlite-vec==0.0.1a37 + protobuf==5.29.2 + py_rust_stemmers==0.1.3 + PyYAML==6.0.2 + requests==2.32.3 + sqlite-vec==0.1.6 + sympy==1.13.3 + tokenizers==0.21.0 + tqdm==4.67.1 + typing_extensions==4.12.2 + urllib3==2.3.0 + watchdog==6.0.0
• server.py
  • @@ -1,39 +1,38 @@
    import mimetypes import os import pathlib - import socketserver import sqlite3 import sys import threading - import traceback - from xmlrpc.server import SimpleXMLRPCDispatcher, SimpleXMLRPCRequestHandler - import pillow_avif + import numpy as np import sqlite_vec from watchdog.observers import Observer from watchdog.events import FileSystemEventHandler, FileOpenedEvent import model + from unixsocket import UnixStreamXMLRPCServer + DIM = 768 + print("Connecting to DB") - con = sqlite3.connect("index.db", check_same_thread=False) + con = sqlite3.connect("search.db") # , check_same_thread=False) + con.execute("PRAGMA journal_mode=WAL") con.enable_load_extension(True) sqlite_vec.load(con) + con.enable_load_extension(False) cur = con.cursor() cur.execute( - "CREATE TABLE IF NOT EXISTS idx (id INTEGER PRIMARY KEY, parent INTEGER, time INTEGER, path TEXT)" - ) - cur.execute( - "CREATE VIRTUAL TABLE IF NOT EXISTS emb USING vec0(id INTEGER PRIMARY KEY, embedding float[1024] distance_metric=cosine)" + f"CREATE VIRTUAL TABLE IF NOT EXISTS idx USING vec0(ino INTEGER PRIMARY KEY, emb float[{DIM}], time INTEGER, parent INTEGER, path TEXT)", ) con.commit() lock = threading.Lock() def get_parent(path): + # Get inode of parent if path in watchdirs: - parent = "/" - else: - parent = pathlib.Path(path).parent + return 0 + parent = pathlib.Path(path).parent return os.stat(parent).st_ino
  • @@ -41,7 +40,7 @@ class EventHandler(FileSystemEventHandler):
    def dispatch(self, event): if not isinstance(event, FileOpenedEvent): with lock: - print(event) + # print(event) super().dispatch(event) def on_created(self, event):
  • @@ -52,17 +51,16 @@ if not event.is_directory:
    self.on_created(event) def on_deleted(self, event): - res = cur.execute("SELECT id FROM idx WHERE path = ?", (event.src_path,)) - ids = res.fetchall() - if len(ids) == 1: - unindex(ids[0][0]) + res = cur.execute("SELECT ino FROM idx WHERE path = ?", (event.src_path,)) + unindex(res.fetchone()[0]) def on_moved(self, event): # inode doesn't change after move + print("Moving", event.src_path, event.dest_path) s = os.stat(event.dest_path) cur.execute( - "INSERT OR REPLACE INTO idx VALUES (?, ?, ?, ?)", - (s.st_ino, get_parent(event.dest_path), s.st_mtime, event.dest_path), + "UPDATE idx SET time = ?, parent = ?, path = ? WHERE ino = ?", + (s.st_mtime_ns, get_parent(event.dest_path), event.dest_path, s.st_ino), ) cur.execute( "UPDATE idx SET path = replace(path, ?, ?)",
  • @@ -80,91 +78,59 @@ print("Indexing", path, parent)
    s = os.stat(path) if os.path.isfile(path): res = cur.execute( - "SELECT time, parent, path FROM idx WHERE id = ?", (s.st_ino,) + "SELECT time, parent, path FROM idx WHERE ino = ?", (s.st_ino,) ) - db_vals = res.fetchall() - if len(db_vals) == 1 and (s.st_mtime, parent, path) == db_vals[0]: - # Already in DB, unmodified - return + db_vals = res.fetchone() - if ( - len(db_vals) == 0 - or s.st_mtime != db_vals[0][1] - or len( - cur.execute("SELECT 1 FROM emb WHERE id = ?", (s.st_ino,)).fetchall() - ) - == 0 - ): - # Modified or not in emb - emb = None + if db_vals is None or db_vals[0] != s.st_mtime_ns or db_vals[1] != parent: + # Not in DB or modified + # Probably faster to query DB first instead of guessing mimetype first type = mimetypes.guess_type(path)[0] - if isinstance(type, str): - try: - if type.startswith("audio"): - emb = model.embed_audio(path) - elif type.startswith("image"): - emb = model.embed_image(path) - elif type.startswith("video") and os.path.getsize(path) < 2**25: - emb = model.embed_video(path) - except: - print(traceback.format_exc()) - - if emb is None: - # Might be in index but no longer valid - unindex(s.st_ino) + if not type.startswith("image"): + # Only support image embeddings for now return - # sqlite-vec doesn't support INSERT OR REPLACE? - cur.execute("DELETE FROM emb WHERE id = ?", (s.st_ino,)) - cur.execute("INSERT INTO emb VALUES (?, ?)", (s.st_ino, emb.cpu().numpy())) + emb = model.embed_image(path) - cur.execute( - "INSERT OR REPLACE INTO idx VALUES (?, ?, ?, ?)", - (s.st_ino, parent, s.st_mtime, path), - ) - con.commit() + # sqlite-vec doesn't support INSERT OR REPLACE and UPSERT + # https://github.com/asg017/sqlite-vec/issues/127 + cur.execute( + "INSERT OR REPLACE INTO idx VALUES (?, ?, ?, ?, ?)", + (s.st_ino, emb, s.st_mtime_ns, parent, path), + ) + con.commit() - if os.path.isdir(path): - if parent: - children = os.listdir(path) - else: - children = watchdirs + elif db_vals[2] != path: + # Moved + cur.execute("UPDATE idx SET path = ? WHERE ino = ?", (path, s.st_ino)) + con.commit() - # Find and unindex dead children - children_id = set( - os.stat(os.path.join(path, child)).st_ino for child in children + elif os.path.isdir(path): + cur.execute( + "INSERT OR REPLACE INTO idx VALUES (?, ?, ?, ?, ?)", + (s.st_ino, np.zeros((DIM,), dtype=np.float32), s.st_mtime_ns, parent, path), ) - res = cur.execute("SELECT id FROM idx WHERE parent = ?", (s.st_ino,)) - db_children_id = res.fetchall() - for db_child_id in db_children_id: - if db_child_id[0] not in children_id: - # Don't unemb because might be a move - unindex(db_child_id[0], False) + con.commit() - # Index live children - for child in children: - if not parent: - observer.schedule(event_handler, child, recursive=True) + for child in os.listdir(path): index(os.path.join(path, child), s.st_ino) - def unindex(id, unemb=True): - print("Unindexing", id) - res = cur.execute("SELECT id FROM idx WHERE parent = ?", (id,)) - db_children_id = res.fetchall() - for db_child_id in db_children_id: - unindex(db_child_id[0]) - cur.execute("DELETE FROM idx WHERE id = ?", (id,)) - if unemb: - cur.execute("DELETE FROM emb WHERE id = ?", (id,)) + def unindex(ino): + print("Unindexing", ino) + res = cur.execute("SELECT id FROM idx WHERE parent = ?", (ino,)) + for db_child in res.fetchall(): + unindex(db_child[0]) + cur.execute("DELETE FROM idx WHERE ino = ?", (ino,)) con.commit() def search(text, limit): + # TODO: Search using image path print("Search", text, limit) - emb = model.embed_text(text).cpu().numpy() + emb = model.embed_text(text) res = cur.execute( - "SELECT idx.path FROM emb JOIN idx ON emb.id = idx.id WHERE embedding MATCH ? AND k = ? ORDER BY distance", + "SELECT path FROM idx WHERE emb MATCH ? AND k = ? ORDER BY distance", (emb, limit), ) return [i[0] for i in res.fetchall()]
  • @@ -175,39 +141,20 @@ watchdirs = set(map(os.path.abspath, sys.argv[1:]))
    observer = Observer() observer.start() event_handler = EventHandler() - with lock: - # Pretend that / is the parent of all indexed dirs - index("/", 0) - # Clean up emb - cur.execute("DELETE FROM emb WHERE id NOT IN (SELECT id FROM idx)") - con.commit() - - - class UnixStreamXMLRPCRequestHandler(SimpleXMLRPCRequestHandler): - disable_nagle_algorithm = False + for wdir in watchdirs: + observer.schedule(event_handler, wdir, recursive=True) - def address_string(self): - return self.client_address + with lock: + for wdir in watchdirs: + # Pretend 0 is parent of watchdirs + index(wdir, 0) - - class UnixStreamXMLRPCServer(socketserver.UnixStreamServer, SimpleXMLRPCDispatcher): - def __init__( - self, - addr, - log_requests=True, - allow_none=True, - encoding=None, - bind_and_activate=True, - use_builtin_types=True, - ): - self.logRequests = log_requests - SimpleXMLRPCDispatcher.__init__(self, allow_none, encoding, use_builtin_types) - socketserver.UnixStreamServer.__init__( - self, - addr, - UnixStreamXMLRPCRequestHandler, - bind_and_activate, - ) + # Remove stale entries + res = cur.execute("SELECT ino, path FROM idx") + for ino, path in res.fetchall(): + if not os.path.exists(path): + cur.execute("DELETE FROM idx WHERE ino = ?", (ino,)) + con.commit() print("Starting RPC server")
• unixsocket.py (new)
  • @@ -0,0 +1,55 @@
    + from http.client import HTTPConnection + import socket + import socketserver + import xmlrpc.client + from xmlrpc.server import SimpleXMLRPCDispatcher, SimpleXMLRPCRequestHandler + + + # https://stackoverflow.com/questions/11729159/use-python-xmlrpclib-with-unix-domain-sockets + + # Client + + + class UnixStreamHTTPConnection(HTTPConnection): + def connect(self): + self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + self.sock.connect(self.host) + + + class UnixStreamTransport(xmlrpc.client.Transport, object): + def __init__(self, socket_path): + self.socket_path = socket_path + super().__init__() + + def make_connection(self, host): + return UnixStreamHTTPConnection(self.socket_path) + + + # Server + + + class UnixStreamXMLRPCRequestHandler(SimpleXMLRPCRequestHandler): + disable_nagle_algorithm = False + + def address_string(self): + return self.client_address + + + class UnixStreamXMLRPCServer(socketserver.UnixStreamServer, SimpleXMLRPCDispatcher): + def __init__( + self, + addr, + log_requests=True, + allow_none=True, + encoding=None, + bind_and_activate=True, + use_builtin_types=True, + ): + self.logRequests = log_requests + SimpleXMLRPCDispatcher.__init__(self, allow_none, encoding, use_builtin_types) + socketserver.UnixStreamServer.__init__( + self, + addr, + UnixStreamXMLRPCRequestHandler, + bind_and_activate, + )