Removed logging file, added reset_dbs method

src/dataset/database_handler.py

@@ -15,10 +15,15 @@ class DatabaseHandler:
         
         self.create_tables()
 
-    def create_tables(self):
+    def create_tables(self, reset: bool = False):
         with sqlite3.connect(self.db_name) as conn:
             cursor = conn.cursor()
             try:
+                if reset:
+                    # Drop the tables if reset is True
+                    cursor.execute("DROP TABLE IF EXISTS entry_database")
+                    cursor.execute("DROP TABLE IF EXISTS chunk_database")
+                
                 # Create entry table
                 query = """
                     CREATE TABLE IF NOT EXISTS entry_database (

src/dataset/update_dataset.py

@@ -2,7 +2,6 @@
 
 from typing import Dict, List
 import numpy as np
-import logging
 from tqdm.auto import tqdm
 import openai
 from datasets import load_dataset
@@ -11,6 +10,10 @@ from .text_splitter import TokenSplitter
 from .pinecone_handler import PineconeHandler
 from .database_handler import DatabaseHandler
 
+import logging
+logger = logging.getLogger(__name__)
+
+
 class ARDUpdater:
     def __init__(
         self, 
@@ -31,33 +34,14 @@ class ARDUpdater:
             max_tokens=max_tokens_per_block
         )
         self.db = DatabaseHandler()
-        self.pinecone_handler = PineconeHandler()
-        
-        ### initialization code ###
-        
-        self.logger = logging.getLogger(__name__)
-        self.logger.setLevel(logging.INFO)
-
-        # File handler
-        file_handler = logging.FileHandler(r'src/dataset/logs/ard_updater.log')
-        file_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
-
-        # Console handler
-        console_handler = logging.StreamHandler()
-        console_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
-
-        self.logger.addHandler(file_handler)
-        self.logger.addHandler(console_handler)
-
-        self.logger.info("ARDUpdater initialized.")
-
+        self.pinecone_db = PineconeHandler()
 
     def update(self, custom_sources: List[str] = ['all']):
         for source in custom_sources:
             self.update_source(source)
         
     def update_source(self, source: str):
-        self.logger.info(f"Updating {source} entries...")
+        logger.info(f"Updating {source} entries...")
 
         iterable_data = load_dataset('StampyAI/alignment-research-dataset', source, split='train', streaming=True)
         iterable_data = iterable_data.map(self.preprocess)
@@ -66,18 +50,18 @@ class ARDUpdater:
         
         for entry in tqdm(iterable_data):
             try:
-                self.pinecone_handler.delete_entry(entry['id'])
+                self.pinecone_db.delete_entry(entry['id'])
 
                 signature = f"Title: {entry['title']}, Authors: {get_authors_str(entry['authors'])}"
                 chunks = self.token_splitter.split(entry['text'], signature)
                 embeddings = self.get_embeddings(chunks)
                 
                 self.db.upsert_chunks(entry['id'], chunks)
-                self.pinecone_handler.insert_entry(entry, chunks, embeddings)
+                self.pinecone_db.insert_entry(entry, chunks, embeddings)
             except Exception as e:
-                self.logger.error(f"An error occurred while updating source {source}: {str(e)}", exc_info=True)
+                logger.error(f"An error occurred while updating source {source}: {str(e)}", exc_info=True)
             
-        self.logger.info(f"Successfully updated {source} entries.")
+        logger.info(f"Successfully updated {source} entries.")
         
     def preprocess(self, entry):
         try:
@@ -93,7 +77,7 @@ class ARDUpdater:
                 'authors': entry['authors']
             }
         except ValueError as e:
-            self.logger.error(f"Entry validation failed: {str(e)}", exc_info=True)
+            logger.error(f"Entry validation failed: {str(e)}", exc_info=True)
             return None
 
     def validate_entry(self, entry: Dict[str, str | list], len_lower_limit: int = 0):
@@ -127,6 +111,10 @@ class ARDUpdater:
         
         return embeddings
 
+    def reset_dbs(self):
+        self.db.create_tables(True)
+        self.pinecone_db.create_index(True)
+
 
 ##### Helper functions #####
 

src/main.py

@@ -20,7 +20,7 @@ def update_database_and_pinecone():
         min_tokens_per_block=200,
         max_tokens_per_block=300,
     )
-    updater.update(['yudkowsky_blog'])
+    updater.update(['gwern_blog'])
 
 
 def main():