Rename search_types to search_type to standardize to singular naming

Using singular names for other directories in application already
 - processor instead of processors
 - interface instead of interfaces

search_type/asymmetric.py

@@ -0,0 +1,177 @@
+#!/usr/bin/env python
+
+import json
+from sentence_transformers import SentenceTransformer, CrossEncoder, util
+import time
+import gzip
+import os
+import sys
+import re
+import torch
+import argparse
+import pathlib
+
+
+def initialize_model():
+    "Initialize model for assymetric semantic search. That is, where query smaller than results"
+    bi_encoder = SentenceTransformer('sentence-transformers/msmarco-MiniLM-L-6-v3')  # The bi-encoder encodes all entries to use for semantic search
+    top_k = 100                                                                      # Number of entries we want to retrieve with the bi-encoder
+    cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')             # The cross-encoder re-ranks the results to improve quality
+    return bi_encoder, cross_encoder, top_k
+
+
+def extract_entries(notesfile, verbose=False):
+    "Load entries from compressed jsonl"
+    entries = []
+    with gzip.open(get_absolute_path(notesfile), 'rt', encoding='utf8') as jsonl:
+        for line in jsonl:
+            note = json.loads(line.strip())
+
+            # Ignore title notes i.e notes with just headings and empty body
+            if not "Body" in note or note["Body"].strip() == "":
+                continue
+
+            note_string = f'{note["Title"]}\t{note["Tags"] if "Tags" in note else ""}\n{note["Body"] if "Body" in note else ""}'
+            entries.extend([note_string])
+
+    if verbose:
+        print(f"Loaded {len(entries)} entries from {notesfile}")
+
+    return entries
+
+
+def compute_embeddings(entries, bi_encoder, embeddings_file, regenerate=False, verbose=False):
+    "Compute (and Save) Embeddings or Load Pre-Computed Embeddings"
+    # Load pre-computed embeddings from file if exists
+    if embeddings_file.exists() and not regenerate:
+        corpus_embeddings = torch.load(get_absolute_path(embeddings_file))
+        if verbose:
+            print(f"Loaded embeddings from {embeddings_file}")
+
+    else:  # Else compute the corpus_embeddings from scratch, which can take a while
+        corpus_embeddings = bi_encoder.encode(entries, convert_to_tensor=True, show_progress_bar=True)
+        torch.save(corpus_embeddings, get_absolute_path(embeddings_file))
+        if verbose:
+            print(f"Computed embeddings and save them to {embeddings_file}")
+
+    return corpus_embeddings
+
+
+def query_notes(raw_query, corpus_embeddings, entries, bi_encoder, cross_encoder, top_k=100):
+    "Search all notes for entries that answer the query"
+    # Separate natural query from explicit required, blocked words filters
+    query = " ".join([word for word in raw_query.split() if not word.startswith("+") and not word.startswith("-")])
+    required_words = set([word[1:].lower() for word in raw_query.split() if word.startswith("+")])
+    blocked_words = set([word[1:].lower() for word in raw_query.split() if word.startswith("-")])
+
+    # Encode the query using the bi-encoder
+    question_embedding = bi_encoder.encode(query, convert_to_tensor=True)
+
+    # Find relevant entries for the query
+    hits = util.semantic_search(question_embedding, corpus_embeddings, top_k=top_k)
+    hits = hits[0]  # Get the hits for the first query
+
+    # Filter results using explicit filters
+    hits = explicit_filter(hits, entries, required_words, blocked_words)
+    if hits is None or len(hits) == 0:
+        return hits
+
+    # Score all retrieved entries using the cross-encoder
+    cross_inp = [[query, entries[hit['corpus_id']]] for hit in hits]
+    cross_scores = cross_encoder.predict(cross_inp)
+
+    # Store cross-encoder scores in results dictionary for ranking
+    for idx in range(len(cross_scores)):
+        hits[idx]['cross-score'] = cross_scores[idx]
+
+    # Order results by cross encoder score followed by biencoder score
+    hits.sort(key=lambda x: x['score'], reverse=True) # sort by biencoder score
+    hits.sort(key=lambda x: x['cross-score'], reverse=True) # sort by cross encoder score
+
+    return hits
+
+
+def explicit_filter(hits, entries, required_words, blocked_words):
+    hits_by_word_set = [(set(word.lower()
+                             for word
+                             in re.split(
+                                 ',|\.| |\]|\[\(|\)|\{|\}',
+                                 entries[hit['corpus_id']])
+                             if word != ""),
+                         hit)
+                        for hit in hits]
+
+    if len(required_words) == 0 and len(blocked_words) == 0:
+        return hits
+    if len(required_words) > 0:
+        return [hit for (words_in_entry, hit) in hits_by_word_set
+                if required_words.intersection(words_in_entry) and not blocked_words.intersection(words_in_entry)]
+    if len(blocked_words) > 0:
+        return [hit for (words_in_entry, hit) in hits_by_word_set
+                if not blocked_words.intersection(words_in_entry)]
+    return hits
+
+
+def render_results(hits, entries, count=5, display_biencoder_results=False):
+    "Render the Results returned by Search for the Query"
+    if display_biencoder_results:
+        # Output of top hits from bi-encoder
+        print("\n-------------------------\n")
+        print(f"Top-{count} Bi-Encoder Retrieval hits")
+        hits = sorted(hits, key=lambda x: x['score'], reverse=True)
+        for hit in hits[0:count]:
+            print(f"Score: {hit['score']:.3f}\n------------\n{entries[hit['corpus_id']]}")
+
+    # Output of top hits from re-ranker
+    print("\n-------------------------\n")
+    print(f"Top-{count} Cross-Encoder Re-ranker hits")
+    hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True)
+    for hit in hits[0:count]:
+        print(f"CrossScore: {hit['cross-score']:.3f}\n-----------------\n{entries[hit['corpus_id']]}")
+
+
+def collate_results(hits, entries, count=5):
+    return [
+        {
+            "Entry": entries[hit['corpus_id']],
+            "Score": f"{hit['cross-score']:.3f}"
+        }
+        for hit
+        in hits[0:count]]
+
+
+def get_absolute_path(filepath):
+    return str(pathlib.Path(filepath).expanduser().absolute())
+
+
+if __name__ == '__main__':
+    # Setup Argument Parser
+    parser = argparse.ArgumentParser(description="Map Org-Mode notes into (compressed) JSONL format")
+    parser.add_argument('--compressed-jsonl', '-j', required=True, type=pathlib.Path, help="Input file for compressed JSONL formatted notes to compute embeddings from")
+    parser.add_argument('--embeddings', '-e', required=True, type=pathlib.Path, help="File to save/load model embeddings to/from")
+    parser.add_argument('--results-count', '-n', default=5, type=int, help="Number of results to render. Default: 5")
+    parser.add_argument('--interactive', action='store_true', default=False, help="Interactive mode allows user to run queries on the model. Default: true")
+    parser.add_argument('--verbose', action='store_true', default=False, help="Show verbose conversion logs. Default: false")
+    args = parser.parse_args()
+
+    # Initialize Model
+    bi_encoder, cross_encoder, top_k = initialize_model()
+
+    # Extract Entries
+    entries = extract_entries(args.compressed_jsonl, args.verbose)
+
+    # Compute or Load Embeddings
+    corpus_embeddings = compute_embeddings(entries, bi_encoder, args.embeddings, args.verbose)
+
+    # Run User Queries on Entries in Interactive Mode
+    while args.interactive:
+        # get query from user
+        user_query = input("Enter your query: ")
+        if user_query == "exit":
+            exit(0)
+
+        # query notes
+        hits = query_notes(user_query, corpus_embeddings, entries, bi_encoder, cross_encoder, top_k)
+
+        # render results
+        render_results(hits, entries, count=args.results_count)

search_type/image-search.py

@@ -0,0 +1,112 @@
+from sentence_transformers import SentenceTransformer, util
+from PIL import Image
+import torch
+import argparse
+import pathlib
+import copy
+
+
+def initialize_model():
+    # Initialize Model
+    torch.set_num_threads(4)
+    top_k = 3
+    model = SentenceTransformer('clip-ViT-B-32')  #Load the CLIP model
+    return model, top_k
+
+
+def extract_entries(image_directory, verbose=False):
+    image_names = list(image_directory.glob('*.jpg'))
+    if verbose:
+        print(f'Found {len(image_names)} images in {image_directory}')
+    return image_names
+
+
+def compute_embeddings(image_names, model, embeddings_file, verbose=False):
+    "Compute (and Save) Embeddings or Load Pre-Computed Embeddings"
+
+    # Load pre-computed embeddings from file if exists
+    if embeddings_file.exists():
+        image_embeddings = torch.load(embeddings_file)
+        if verbose:
+            print(f"Loaded pre-computed embeddings from {embeddings_file}")
+
+    else:  # Else compute the image_embeddings from scratch, which can take a while
+        images = []
+        if verbose:
+            print(f"Loading the {len(image_names)} images into memory")
+        for image_name in image_names:
+            images.append(copy.deepcopy(Image.open(image_name)))
+
+        if len(images) > 0:
+            image_embeddings = model.encode(images, batch_size=128, convert_to_tensor=True, show_progress_bar=True)
+            torch.save(image_embeddings, embeddings_file)
+            if verbose:
+                print(f"Saved computed embeddings to {embeddings_file}")
+
+    return image_embeddings
+
+
+def search(query, image_embeddings, model, count=3, verbose=False):
+    # Set query to image content if query is a filepath
+    if pathlib.Path(query).expanduser().is_file():
+        query_imagepath = pathlib.Path(query).expanduser().resolve(strict=True)
+        query = copy.deepcopy(Image.open(query_imagepath))
+        if verbose:
+            print(f"Find Images similar to Image at {query_imagepath}")
+    else:
+        print(f"Find Images by Text: {query}")
+
+    # Now we encode the query (which can either be an image or a text string)
+    query_embedding = model.encode([query], convert_to_tensor=True, show_progress_bar=False)
+
+    # Then, we use the util.semantic_search function, which computes the cosine-similarity
+    # between the query embedding and all image embeddings.
+    # It then returns the top_k highest ranked images, which we output
+    hits = util.semantic_search(query_embedding, image_embeddings, top_k=count)[0]
+
+    return hits
+
+
+def render_results(hits, image_names, image_directory, count):
+    for hit in hits[:count]:
+        print(image_names[hit['corpus_id']])
+        image_path = image_directory.joinpath(image_names[hit['corpus_id']])
+        with Image.open(image_path) as img:
+            img.show()
+
+
+if __name__ == '__main__':
+    # Setup Argument Parser
+    parser = argparse.ArgumentParser(description="Semantic Search on Images")
+    parser.add_argument('--image-directory', '-i', required=True, type=pathlib.Path, help="Image directory to query")
+    parser.add_argument('--embeddings-file', '-e', default='embeddings.pt', type=pathlib.Path, help="File to save/load model embeddings to/from. Default: ./embeddings.pt")
+    parser.add_argument('--results-count', '-n', default=5, type=int, help="Number of results to render. Default: 5")
+    parser.add_argument('--interactive', action='store_true', default=False, help="Interactive mode allows user to run queries on the model. Default: true")
+    parser.add_argument('--verbose', action='store_true', default=False, help="Show verbose conversion logs. Default: false")
+    args = parser.parse_args()
+
+    # Resolve file, directory paths in args to absolute paths
+    embeddings_file = args.embeddings_file.expanduser().resolve()
+    image_directory = args.image_directory.expanduser().resolve(strict=True)
+
+    # Initialize Model
+    model, count = initialize_model()
+
+    # Extract Entries
+    image_names = extract_entries(image_directory, args.verbose)
+
+    # Compute or Load Embeddings
+    image_embeddings = compute_embeddings(image_names, model, embeddings_file, args.verbose)
+
+    # Run User Queries on Entries in Interactive Mode
+    while args.interactive:
+        # get query from user
+        user_query = input("Enter your query: ")
+        if user_query == "exit":
+            exit(0)
+
+        # query notes
+        hits = search(user_query, image_embeddings, model, args.results_count, args.verbose)
+
+        # render results
+        render_results(hits, image_names, image_directory, count=args.results_count)

search_type/symmetric.py

@@ -0,0 +1,97 @@
+import pandas as pd
+import faiss
+import numpy as np
+
+from sentence_transformers import SentenceTransformer
+
+import argparse
+import os
+
+def create_index(
+        model,
+        dataset_path,
+        index_path,
+        column_name,
+        recreate):
+    # Load Dataset
+    dataset = pd.read_csv(dataset_path)
+
+    # Clean Dataset
+    dataset = dataset.dropna()
+    dataset[column_name] = dataset[column_name].str.strip()
+
+    # Create Index or Load it if it already exists
+    if os.path.exists(index_path) and not recreate:
+        index = faiss.read_index(index_path)
+    else:
+        # Create Embedding Vectors of Documents
+        embeddings = model.encode(dataset[column_name].to_list(), show_progress_bar=True)
+        embeddings = np.array([embedding for embedding in embeddings]).astype("float32")
+
+        index = faiss.IndexIDMap(
+            faiss.IndexFlatL2(
+                embeddings.shape[1]))
+
+        index.add_with_ids(embeddings, dataset.index.values)
+
+        faiss.write_index(index, index_path)
+
+    return index, dataset
+
+
+def resolve_column(dataset, Id, column):
+    return [list(dataset[dataset.index == idx][column]) for idx in Id[0]]
+
+
+def vector_search(query, index, dataset, column_name, num_results=10):
+    query_vector = np.array(query).astype("float32")
+    D, Id = index.search(query_vector, k=num_results)
+
+    return zip(D[0], Id[0], resolve_column(dataset, Id, column_name))
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description="Find most suitable match based on users exclude, include preferences")
+    parser.add_argument('positives', type=str, help="Terms to find closest match to")
+    parser.add_argument('--negatives', '-n', type=str, help="Terms to find farthest match from")
+
+    parser.add_argument('--recreate', action='store_true', default=False, help="Recreate index at index_path from dataset at dataset path")
+    parser.add_argument('--index', type=str, default="./.faiss_index", help="Path to index for storing vector embeddings")
+    parser.add_argument('--dataset', type=str, default="./.dataset", help="Path to dataset to generate index from")
+    parser.add_argument('--column', type=str, default="DATA", help="Name of dataset column to index")
+    parser.add_argument('--num_results', type=int, default=10, help="Number of most suitable matches to show")
+    parser.add_argument('--model_name', type=str, default='paraphrase-distilroberta-base-v1', help="Specify name of the SentenceTransformer model to use for encoding")
+    args = parser.parse_args()
+
+    model = SentenceTransformer(args.model_name)
+
+    if args.positives and not args.negatives:
+        # Get index, create it from dataset if doesn't exist
+        index, dataset = create_index(model, args.dataset, args.index, args.column, args.recreate)
+
+        # Create vector to represent user's stated positive preference
+        preference_vector = model.encode([args.positives])
+
+        # Find and display most suitable matches for users preferences in the dataset
+        results = vector_search(preference_vector, index, dataset, args.column, args.num_results)
+
+        print("Most Suitable Matches:")
+        for similarity, id_, data in results:
+            print(f"Id: {id_}\nSimilarity: {similarity}\n{args.column}: {data[0]}")
+
+    elif args.positives and args.negatives:
+        # Get index, create it from dataset if doesn't exist
+        index, dataset = create_index(model, args.dataset, args.index, args.column, args.recreate)
+
+        # Create vector to represent user's stated preference
+        positives_vector = np.array(model.encode([args.positives])).astype("float32")
+        negatives_vector = np.array(model.encode([args.negatives])).astype("float32")
+
+        # preference_vector = np.mean([positives_vector, -1 * negatives_vector], axis=0)
+        preference_vector = np.add(positives_vector, -1 * negatives_vector)
+
+        # Find and display most suitable matches for users preferences in the dataset
+        results = vector_search(preference_vector, index, dataset, args.column, args.num_results)
+
+        print("Most Suitable Matches:")
+        for similarity, id_, data in results:
+            print(f"Id: {id_}\nSimilarity: {similarity}\n{args.column}: {data[0]}")