Spaces:
Sleeping
Sleeping
| """ | |
| database.py | |
| =========== | |
| SQLite mirror of the metadata CSV, used only for fast searching/filtering in | |
| the Dataset Explorer tab. The CSV (see ``metadata.py``) remains the master | |
| export — this database is a derived, disposable convenience index. | |
| Why "download, modify, re-upload" instead of querying ``hf://`` directly? | |
| -------------------------------------------------------------------------- | |
| ``sqlite3`` needs a real local file handle with support for locking/seeking, | |
| which object storage / fsspec paths don't reliably provide. So each mutating | |
| operation: | |
| 1. Pulls the current DB file from the bucket to a local scratch path | |
| (skipped if using a mounted volume, where the "local" path *is* already | |
| inside the bucket, i.e. zero-copy). | |
| 2. Opens it with ``sqlite3``, applies the change. | |
| 3. Immediately pushes the updated file back to the bucket. | |
| All of this is wrapped in a process-wide lock so concurrent uploads cannot | |
| interleave and corrupt the database file. | |
| """ | |
| from __future__ import annotations | |
| import logging | |
| import os | |
| import sqlite3 | |
| import threading | |
| from contextlib import contextmanager | |
| from typing import Dict, Iterator, List | |
| from bucket import BucketStorage | |
| from utils import DATABASE_PATH, METADATA_COLUMNS | |
| logger = logging.getLogger("speech_collector.database") | |
| _db_lock = threading.RLock() | |
| _CREATE_TABLE_SQL = f""" | |
| CREATE TABLE IF NOT EXISTS samples ( | |
| {", ".join(f'{col} TEXT' for col in METADATA_COLUMNS if col not in ("sample_id", "age", "audio_duration", "sample_rate"))}, | |
| sample_id TEXT PRIMARY KEY, | |
| age INTEGER, | |
| audio_duration REAL, | |
| sample_rate INTEGER | |
| ) | |
| """ | |
| class SpeechDatabase: | |
| """Thread-safe SQLite mirror living inside the bucket.""" | |
| def __init__(self, bucket: BucketStorage, local_scratch_path: str = "/tmp/speech_dataset_local.db") -> None: | |
| self.bucket = bucket | |
| self.local_scratch_path = local_scratch_path | |
| self._ensure_db_exists() | |
| def _ensure_db_exists(self) -> None: | |
| with _db_lock: | |
| if self.bucket.exists(DATABASE_PATH): | |
| self.bucket.download_to_local(DATABASE_PATH, self.local_scratch_path) | |
| else: | |
| logger.info("speech_dataset.db not found in bucket — creating a new one.") | |
| if os.path.exists(self.local_scratch_path): | |
| os.remove(self.local_scratch_path) | |
| conn = sqlite3.connect(self.local_scratch_path) | |
| try: | |
| conn.execute(_CREATE_TABLE_SQL) | |
| conn.commit() | |
| finally: | |
| conn.close() | |
| self.bucket.sync_local_to_bucket(self.local_scratch_path, DATABASE_PATH) | |
| def _connect_synced(self) -> Iterator[sqlite3.Connection]: | |
| """Pull latest DB from bucket, yield a connection, push changes back.""" | |
| with _db_lock: | |
| if self.bucket.exists(DATABASE_PATH): | |
| self.bucket.download_to_local(DATABASE_PATH, self.local_scratch_path) | |
| conn = sqlite3.connect(self.local_scratch_path) | |
| try: | |
| yield conn | |
| conn.commit() | |
| finally: | |
| conn.close() | |
| self.bucket.sync_local_to_bucket(self.local_scratch_path, DATABASE_PATH) | |
| def insert_sample(self, row: Dict[str, str]) -> None: | |
| columns = METADATA_COLUMNS | |
| placeholders = ", ".join("?" for _ in columns) | |
| values = [row.get(col, "") for col in columns] | |
| sql = f"INSERT OR REPLACE INTO samples ({', '.join(columns)}) VALUES ({placeholders})" | |
| try: | |
| with self._connect_synced() as conn: | |
| conn.execute(sql, values) | |
| logger.info("Mirrored sample %s into SQLite database.", row.get("sample_id")) | |
| except Exception as exc: | |
| # The CSV row is already safely written by this point; a DB | |
| # mirroring failure must not be treated as a fatal upload error. | |
| logger.error("Failed to mirror sample into SQLite (non-fatal): %s", exc) | |
| def fetch_all(self) -> List[tuple]: | |
| try: | |
| with self._connect_synced() as conn: | |
| cursor = conn.execute( | |
| "SELECT sample_id, name, gender, age, region, state, country, " | |
| "language, dialect, prompt_text, transcript, audio_file, " | |
| "audio_duration, timestamp FROM samples ORDER BY sample_id DESC" | |
| ) | |
| return cursor.fetchall() | |
| except Exception as exc: | |
| logger.error("Failed to fetch samples from SQLite: %s", exc) | |
| return [] | |
| def aggregate_stats(self) -> Dict: | |
| try: | |
| with self._connect_synced() as conn: | |
| total_samples, total_speakers, total_duration = conn.execute( | |
| "SELECT COUNT(*), COUNT(DISTINCT name), COALESCE(SUM(audio_duration), 0) FROM samples" | |
| ).fetchone() | |
| def group_counts(column: str) -> List[tuple]: | |
| return conn.execute( | |
| f"SELECT {column}, COUNT(*) FROM samples " | |
| f"WHERE {column} IS NOT NULL AND {column} != '' " | |
| f"GROUP BY {column} ORDER BY COUNT(*) DESC" | |
| ).fetchall() | |
| latest_uploads = conn.execute( | |
| "SELECT sample_id, name, language, timestamp FROM samples " | |
| "ORDER BY sample_id DESC LIMIT 10" | |
| ).fetchall() | |
| age_stats = conn.execute( | |
| "SELECT MIN(age), MAX(age), AVG(age) FROM samples WHERE age IS NOT NULL" | |
| ).fetchone() | |
| return { | |
| "total_samples": total_samples or 0, | |
| "total_speakers": total_speakers or 0, | |
| "total_duration_seconds": total_duration or 0.0, | |
| "gender_distribution": group_counts("gender"), | |
| "language_distribution": group_counts("language"), | |
| "dialect_distribution": group_counts("dialect"), | |
| "latest_uploads": latest_uploads, | |
| "age_min": age_stats[0] if age_stats else None, | |
| "age_max": age_stats[1] if age_stats else None, | |
| "age_avg": round(age_stats[2], 1) if age_stats and age_stats[2] else None, | |
| } | |
| except Exception as exc: | |
| logger.error("Failed to aggregate stats from SQLite: %s", exc) | |
| return { | |
| "total_samples": 0, "total_speakers": 0, "total_duration_seconds": 0.0, | |
| "gender_distribution": [], "language_distribution": [], "dialect_distribution": [], | |
| "latest_uploads": [], "age_min": None, "age_max": None, "age_avg": None, | |
| } | |