import sqlite3 def initialize_database(): """Initialize the SQLite database and create the 'documents' table if it doesn't exist.""" # Connect to the SQLite database (or create it if it doesn't exist) conn = sqlite3.connect('dataset.db') cursor = conn.cursor() # Create the 'documents' table if it doesn't exist cursor.execute(''' CREATE TABLE IF NOT EXISTS documents ( id INTEGER PRIMARY KEY AUTOINCREMENT, text TEXT NOT NULL, topics TEXT ) ''') # Commit changes and close the connection conn.commit() conn.close() def save_to_db(chunks, topics=None): """Save chunks to the SQLite database.""" # Ensure the database and table are initialized initialize_database() # Connect to the database conn = sqlite3.connect('dataset.db') cursor = conn.cursor() # Insert chunks into the database for chunk in chunks: cursor.execute('INSERT INTO documents (text, topics) VALUES (?, ?)', (chunk, topics)) # Commit changes and close the connection conn.commit() conn.close()