stark_qa.retrieval

stark_qa.retrieval.dataset

class stark_qa.retrieval.dataset.STaRKDataset(name, root=None, human_generated_eval=False)

Bases: object

get_idx_split(test_ratio=1.0)

Return the indices of train/val/test split in a dictionary.

Parameters:

test_ratio (float) – Ratio of test data to include.

Returns:

Dictionary with split indices for train, val, and test sets.

Return type:

dict

get_query_by_qid(q_id)

Return the query by query id.

Parameters:

q_id (int) – Query id.

Returns:

Query string.

Return type:

str

get_subset(split)

Return a subset of the dataset.

Parameters:

split (str) – Split type (‘train’, ‘val’, ‘test’).

Returns:

Subset of the dataset.

Return type:

STaRKDataset