Auto-format code with isort and black

benchmarks/retrieval/retrieve.py

@@ -33,7 +33,7 @@ def main():
     parser.add(
         "--logs-dir",
         default=None,
-        help="Path where to output predictions and metrics. Optional, since metrics are also printed to console."
+        help="Path where to output predictions and metrics. Optional, since metrics are also printed to console.",
     )
     parser.add("--max-instances", default=None, type=int, help="Maximum number of instances to process.")
 
@@ -72,9 +72,7 @@ def main():
             # the retrived documents. The key of the score varies depending on the underlying retriever. If there's no
             # score, we use 1/(doc_idx+1) since it preserves the order of the documents.
             score = doc.metadata.get("score", doc.metadata.get("relevance_score", 1 / (doc_idx + 1)))
-            retrieved_docs.append(
-                ScoredDoc(query_id=query_id, doc_id=doc.metadata["file_path"], score=score)
-            )
+            retrieved_docs.append(ScoredDoc(query_id=query_id, doc_id=doc.metadata["file_path"], score=score))
             # Update the output dictionary with the retrieved documents.
             item["retrieved"].append({"file_path": doc.metadata["file_path"], "score": score})
 

benchmarks/retrieval/retrieve_kaggle.py

@@ -40,7 +40,9 @@ def main():
 
         retrieved = retriever.invoke(item["question"])
         # Sort by score in descending order.
-        retrieved = sorted(retrieved, key=lambda doc: doc.metadata.get("score", doc.metadata.get("relevance_score")), reverse=True)
+        retrieved = sorted(
+            retrieved, key=lambda doc: doc.metadata.get("score", doc.metadata.get("relevance_score")), reverse=True
+        )
         # Keep top 3, since the Kaggle competition only evaluates the top 3.
         retrieved = retrieved[:3]
         retrieved_filenames = [doc.metadata["file_path"] for doc in retrieved]

sage/index.py

@@ -42,7 +42,6 @@ def main():
     if args.embedding_provider == "marqo" and args.vector_store_provider != "marqo":
         parser.error("When using the marqo embedder, the vector store type must also be marqo.")
 
-    
     ######################
     # Step 1: Embeddings #
     ######################

sage/vector_store.py

@@ -1,13 +1,13 @@
 """Vector store abstraction and implementations."""
 
 import logging
-import nltk
 import os
 from abc import ABC, abstractmethod
 from functools import cached_property
 from typing import Dict, Generator, List, Optional, Tuple
 
 import marqo
+import nltk
 from langchain_community.retrievers import PineconeHybridSearchRetriever
 from langchain_community.vectorstores import Marqo
 from langchain_community.vectorstores import Pinecone as LangChainPinecone
@@ -22,13 +22,15 @@ from sage.data_manager import DataManager
 
 Vector = Tuple[Dict, List[float]]  # (metadata, embedding)
 
+
 def is_punkt_downloaded():
     try:
-        find('tokenizers/punkt_tab')
+        find("tokenizers/punkt_tab")
         return True
     except LookupError:
         return False
 
+
 class VectorStore(ABC):
     """Abstract class for a vector store."""
 
@@ -83,7 +85,7 @@ class PineconeVectorStore(VectorStore):
                 else:
                     print("punkt is not downloaded")
                     # Optionally download it
-                    nltk.download('punkt_tab')
+                    nltk.download("punkt_tab")
                 self.bm25_encoder = BM25Encoder()
                 self.bm25_encoder.load(path=bm25_cache)
             else: