Spaces:
Sleeping
Sleeping
File size: 6,186 Bytes
4629ca3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 | """Data preprocessing and grounding validation pipeline."""
import hashlib
import logging
from collections.abc import Iterable
from pathlib import Path
from typing import Any
from pydantic import BaseModel, ConfigDict
from tqdm import tqdm
from ask_the_database.environment.schema_manager import SchemaManager
from ask_the_database.environment.sqlite_env import SQLiteEnvironment
logger = logging.getLogger(__name__)
class ProcessedSample(BaseModel):
"""Normalized, grounded question-SQL sample adhering to canonical schema."""
model_config = ConfigDict(frozen=True)
instance_id: str
source: str
db_id: str
question: str
gold_sql: str
db_path: str
schema_ddl: str
difficulty: str
gold_row_count: int
gold_execution_hash: str
class DataPreprocessor:
"""Validates, grounds, and normalizes raw dataset records into canonical JSONL formats."""
def __init__(
self,
databases_dir: str = "data/databases",
processed_dir: str = "data/processed",
timeout_seconds: float = 5.0,
) -> None:
self.databases_dir = Path(databases_dir)
self.processed_dir = Path(processed_dir)
self.processed_dir.mkdir(parents=True, exist_ok=True)
self.timeout_seconds = timeout_seconds
self._schema_cache: dict[str, str] = {}
def get_database_path(self, db_id: str) -> Path | None:
"""Find the physical .sqlite file for a given database ID."""
candidate = self.databases_dir / db_id / f"{db_id}.sqlite"
if candidate.exists():
return candidate
matches = list(self.databases_dir.glob(f"**/{db_id}.sqlite"))
if matches:
return matches[0]
return None
def get_schema_ddl(self, db_path: Path) -> str:
"""Fetch and cache schema DDL string for a database."""
path_str = str(db_path)
if path_str not in self._schema_cache:
manager = SchemaManager(path_str)
self._schema_cache[path_str] = manager.get_schema_summary()
return self._schema_cache[path_str]
@staticmethod
def compute_execution_hash(data: list[tuple[Any, ...]] | None) -> str:
"""Generate a deterministic SHA-256 hash from query result rows."""
if not data:
return hashlib.sha256(b"EMPTY_RESULT").hexdigest()
normalized = repr(sorted(str(row) for row in data))
return hashlib.sha256(normalized.encode("utf-8")).hexdigest()
def process_record(
self,
record: dict[str, Any],
source: str,
split: str,
index: int,
) -> ProcessedSample | None:
"""Process a single record, verify its gold SQL executes, and return ProcessedSample."""
db_id = record.get("db_id")
if not db_id:
logger.warning("Record %s missing db_id, skipping.", index)
return None
# Extract SQL and question across Spider and SkyRL naming conventions
gold_sql = record.get("gold_sql") or record.get("query") or record.get("sql")
question = record.get("question")
difficulty = record.get("difficulty", "medium")
if not gold_sql or not question:
logger.warning("Record %s missing question or SQL, skipping.", index)
return None
db_path = self.get_database_path(db_id)
if not db_path:
logger.warning(
"Database file for %s not found in %s, skipping.",
db_id,
self.databases_dir,
)
return None
# Schema extraction
try:
schema_ddl = self.get_schema_ddl(db_path)
except Exception as e:
logger.warning("Failed extracting schema for %s: %s", db_id, e)
return None
# Execute gold SQL in sandbox to verify execution correctness
env = SQLiteEnvironment(
str(db_path),
timeout_seconds=self.timeout_seconds,
enforce_ast_safety=True,
)
res = env.execute(gold_sql)
if not res.success:
logger.warning(
"Gold SQL failed execution for db_id=%s [sample %s]: %s (Query: %s)",
db_id,
index,
res.error_message,
gold_sql,
)
return None
instance_id = f"{source}_{split}_{index:05d}"
exec_hash = self.compute_execution_hash(res.data)
# Store relative db_path for portability
try:
rel_db_path = str(db_path.relative_to(Path.cwd()))
except ValueError:
rel_db_path = str(db_path)
return ProcessedSample(
instance_id=instance_id,
source=source,
db_id=db_id,
question=question.strip(),
gold_sql=gold_sql.strip(),
db_path=rel_db_path,
schema_ddl=schema_ddl,
difficulty=str(difficulty).lower(),
gold_row_count=res.row_count,
gold_execution_hash=exec_hash,
)
def process_split(
self,
records: Iterable[dict[str, Any]],
source: str,
split: str,
output_file: Path | str,
max_samples: int | None = None,
) -> int:
"""Process an entire split and write valid samples to a JSONL file."""
out_path = Path(output_file)
out_path.parent.mkdir(parents=True, exist_ok=True)
valid_count = 0
total_seen = 0
with open(out_path, "w", encoding="utf-8") as f:
for idx, record in enumerate(tqdm(records, desc=f"Processing {source}:{split}")):
if max_samples and valid_count >= max_samples:
break
total_seen += 1
sample = self.process_record(record, source=source, split=split, index=idx)
if sample is not None:
f.write(sample.model_dump_json() + "\n")
valid_count += 1
logger.info(
"Wrote %d valid samples out of %d seen to %s",
valid_count,
total_seen,
out_path,
)
return valid_count
|