fixed typing issues

StampyAI · Aug 29, 2023 · 1238ac8 · 1238ac8
1 parent a457fa8
commit 1238ac8
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 2 deletions.
diff --git a/align_data/embeddings/finetuning/finetuning_dataset.py b/align_data/embeddings/finetuning/finetuning_dataset.py
@@ -20,7 +20,7 @@
 class FinetuningDataset(IterableDataset):
     def __init__(self, num_batches_per_epoch: int, cache_size: int = 1280):
         self.num_batches_per_epoch = num_batches_per_epoch
-        self.article_cache = deque(maxlen=cache_size)
+        self.article_cache: deque = deque(maxlen=cache_size)
 
         self.text_splitter = ParagraphSentenceUnitTextSplitter()
         self.pinecone_db = PineconeDB()

diff --git a/align_data/embeddings/pinecone/pinecone_db_handler.py b/align_data/embeddings/pinecone/pinecone_db_handler.py
@@ -100,6 +100,10 @@ def query_text(
         **kwargs,
     ) -> List[ScoredVector]:
         query_vector = get_embedding(query)[0]
+        if query_vector is None:
+            print("The query is invalid.")
+            return []
+
         return self.query_vector(
             query=query_vector,
             top_k=top_k,

diff --git a/align_data/embeddings/pinecone/update_pinecone.py b/align_data/embeddings/pinecone/update_pinecone.py
@@ -55,7 +55,7 @@ def update_articles_by_ids(self, hash_ids: List[int], force_update: bool = False
             for batch in self.batch_entries(articles_to_update):
                 self.save_batch(session, batch)
 
-    def save_batch(self, session: Session, batch: List[Tuple[Article, PineconeEntry]]):
+    def save_batch(self, session: Session, batch: List[Tuple[Article, PineconeEntry | None]]):
         try:
             for article, pinecone_entry in batch:
                 if pinecone_entry: