helpers/TextSearch.py

import gc

import pyarrow as pa
import pyarrow.parquet as pq
import faiss
import torch
from open_clip import create_model_from_pretrained, get_tokenizer

# Columns actually used downstream. The metadata parquet also carries a
# 'geometry' column that is ~2GB uncompressed and is never read by this app,
# so it is deliberately left on disk.
METADATA_COLUMNS = ['grid_cell', 'file', 'row_idx']

class SearchSigLIP():

    def __init__(self, index_path, metadata_path):
        # Everything here is loaded lazily, on the first text search.
        #
        # Eagerly loading the faiss index (~1GB), the metadata table and the
        # SigLIP weights (~3.5GB) costs more RAM than a cpu-basic Space has,
        # and it happens at import time - so the container was being OOM-killed
        # before it could serve the map, which does not need any of it.
        self.index_path = index_path
        self.metadata_path = metadata_path
        self._ready = False

    def _ensure_ready(self):
        if self._ready:
            return

        print(f'Loading index from PATH={self.index_path}', flush=True)
        self.init_index()
        print('[DONE]', flush=True)

        print(f'Loading metadata from PATH={self.metadata_path}', flush=True)
        self.metadata = pq.read_table(self.metadata_path, columns=METADATA_COLUMNS)
        print('[DONE]', flush=True)

        self.init_model()
        self._ready = True

    def init_index(self):
        self.cpu_index = faiss.read_index(self.index_path)

        # Only move the index onto a GPU when one is actually usable. On CPU-only
        # hardware faiss is built without the GPU symbols, so StandardGpuResources
        # does not exist at all and we search the CPU index directly.
        if hasattr(faiss, 'StandardGpuResources') and faiss.get_num_gpus() > 0:
            res = faiss.StandardGpuResources()
            cloner_options = faiss.GpuClonerOptions()
            cloner_options.useFloat16LookupTables = True
            self.index = faiss.index_cpu_to_gpu(res, 0, self.cpu_index, cloner_options)
        else:
            print('No GPU available for faiss - searching the CPU index.', flush=True)
            self.index = self.cpu_index

        self.index.nprobe = 32 # Higher = more accurate, slower

    def init_model(self):
        self.model, self.preprocess = create_model_from_pretrained('hf-hub:timm/ViT-SO400M-14-SigLIP-384')
        self.model.eval()
        self.tokenizer = get_tokenizer('hf-hub:timm/ViT-SO400M-14-SigLIP')

        # Only encode_text is ever called here, so the vision tower is dead
        # weight - roughly half the parameters. Dropping it frees over a GB.
        # Guarded: if a future open_clip needs it, keeping the tower only costs
        # memory, whereas a hard failure here would break search entirely.
        try:
            if hasattr(self.model, 'visual'):
                del self.model.visual
                gc.collect()
        except Exception as e:
            print(f'Could not release the vision tower: {e}', flush=True)

    def encode_text(self, text, device=None):
        if device is None:
            device = 'cuda' if torch.cuda.is_available() else 'cpu'
        self.model.to(device)
        with torch.no_grad():
            text = self.tokenizer([text], context_length=self.model.context_length)
            return self.model.encode_text(text.to(device))

    def search_with_grid(self, query_vec, k=5):
        # Prepare query
        if isinstance(query_vec, torch.Tensor):
            query_vec = query_vec.cpu().squeeze().numpy()
            
        query_vec = query_vec.reshape(1, -1).astype('float32')
        faiss.normalize_L2(query_vec)
        
        # Search
        distances, indices = self.index.search(query_vec, k)
        
        # Flatten results
        ids = indices[0]
        scores = distances[0]
        
        # We ignore -1 (which happens if k > total vectors, unlikely here)
        valid_mask = ids != -1
        valid_ids = ids[valid_mask]
        valid_scores = scores[valid_mask]

        if len(valid_ids) == 0:
            return []

        # Direct lookup by integer index, straight out of the Arrow table.
        # Kept in Arrow rather than pandas: as an object-dtype DataFrame the
        # 20M grid_cell strings alone cost well over a gigabyte.
        matches = self.metadata.take(pa.array(valid_ids.astype('int64')))

        results = matches.to_pylist()
        for row, score in zip(results, valid_scores):
            row['score'] = float(score)

        return results

    def faiss(self, text, k=1): # k - number of neighbours

        # 0. Load the index/metadata/model if this is the first search
        self._ensure_ready()

        # 1. Compute query
        q = self.encode_text(text)

        # 2. Find Hits
        results = self.search_with_grid(q, k=k)

        return results