| """Database-based dataset splitter.""" |
|
|
| import random |
| from typing import List, Set, Tuple |
|
|
| from src.t2s.data.spider_loader import SpiderExample |
|
|
|
|
| def split_train_val_by_db( |
| examples: List[SpiderExample], val_ratio: float = 0.10, seed: int = 42 |
| ) -> Tuple[List[SpiderExample], List[SpiderExample], Set[str], Set[str]]: |
| unique_dbs = sorted(list({ex.db_id for ex in examples})) |
|
|
| rng = random.Random(seed) |
| rng.shuffle(unique_dbs) |
|
|
| num_val_dbs = max(1, int(len(unique_dbs) * val_ratio)) |
| val_dbs = set(unique_dbs[:num_val_dbs]) |
| train_dbs = set(unique_dbs[num_val_dbs:]) |
|
|
| train_examples = [ex for ex in examples if ex.db_id in train_dbs] |
| val_examples = [ex for ex in examples if ex.db_id in val_dbs] |
|
|
| return train_examples, val_examples, train_dbs, val_dbs |
|
|
