Draft of feat: allow "negative" text queries

rclip/main.py

@@ -121,10 +121,13 @@ def ensure_index(self, directory: str):
 
     self._db.commit()
 
-  def search(self, query: str, directory: str, top_k: int = 10) -> List[SearchResult]:
+  def search(
+      self, query: str, directory: str, top_k: int = 10,
+      positive_queries: List[str] = [], negative_queries: List[str] = []) -> List[SearchResult]:
     filepaths, features = self._get_features(directory)
 
-    sorted_similarities = self._model.compute_similarities_to_text(features, query)
+    positive_queries = [query] + positive_queries
+    sorted_similarities = self._model.compute_similarities_to_text(features, positive_queries, negative_queries)
 
     filtered_similarities = filter(
       lambda similarity: not self._exclude_dir_regex.match(filepaths[similarity[1]]),
@@ -159,7 +162,7 @@ def main():
   if not args.skip_index:
     rclip.ensure_index(current_directory)
 
-  result = rclip.search(args.query, current_directory, args.top)
+  result = rclip.search(args.query, current_directory, args.top, args.add, args.subtract)
   if args.filepath_only:
     for r in result:
       print(r.filepath)

rclip/model.py

@@ -30,17 +30,24 @@ def compute_image_features(self, images: List[Image.Image]) -> np.ndarray:
 
     return image_features.cpu().numpy()
 
-  def compute_text_features(self, text: str) -> np.ndarray:
+  def compute_text_features(self, text: List[str]) -> np.ndarray:
     with torch.no_grad():
       text_encoded = self._model.encode_text(clip.tokenize(text).to(self._device))
       text_encoded /= text_encoded.norm(dim=-1, keepdim=True)
 
     return text_encoded.cpu().numpy()
 
-  def compute_similarities_to_text(self, item_features: np.ndarray, text: str) -> List[Tuple[float, int]]:
-    text_features = self.compute_text_features(text)
+  def compute_similarities_to_text(
+      self, item_features: np.ndarray,
+      positive_queries: List[str], negative_queries: List[str]) -> List[Tuple[float, int]]:
 
-    similarities = (text_features @ item_features.T).squeeze(0).tolist()
+    positive_features = self.compute_text_features(positive_queries)
+    text_features = np.add.reduce(positive_features)
+    if negative_queries:
+        negative_features = self.compute_text_features(negative_queries)
+        text_features -= np.add.reduce(negative_features)
+
+    similarities = text_features @ item_features.T
     sorted_similarities = sorted(zip(similarities, range(item_features.shape[0])), key=lambda x: x[0], reverse=True)
 
     return sorted_similarities

rclip/utils.py

@@ -48,6 +48,9 @@ def top_arg_type(arg: str) -> int:
 def init_arg_parser() -> argparse.ArgumentParser:
   parser = argparse.ArgumentParser()
   parser.add_argument('query')
+  parser.add_argument('--add', '-a', action='append', default=[], help='queries to add to the "original" query')
+  parser.add_argument('--subtract', '--sub', '-s', action='append', default=[],
+                      help='queries to subtract from the "original" query')
   parser.add_argument('--top', '-t', type=top_arg_type, default=10, help='number of top results to display')
   parser.add_argument('--filepath-only', '-f', action='store_true', default=False, help='outputs only filepaths')
   parser.add_argument(