database.py

import sqlite3
import os

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()

from huggingface_hub import HfApi

def commit_to_huggingface():
    """Commit the dataset.db file to the Hugging Face Space repository."""
    api_token = os.getenv("hf_key")
    api = HfApi(token=api_token)

    # Replace with your Space's repository name
    repo_id = "Danielrahmai1991/dataset_interface"

    # Upload and commit the dataset.db file
    api.upload_file(
        path_or_fileobj="dataset.db",
        path_in_repo="dataset.db",
        repo_id=repo_id,
        repo_type="space"
    )



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()
    commit_to_huggingface()