File size: 790 Bytes
388137e | 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 | """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
|