MatchPrePrintArticles / fetch_positive_samples.py
KNGCRIMSON's picture
app
b5cf002
from src.dataset.GoodDataset import *
import argparse
def main(config):
"""
Main function to process the dataset and save it as a CSV file.
Args:
config: Namespace object containing the script arguments.
"""
# Initialize the dataset
dataset = AugmentedDataset()
positive_samples = dataset.fetch_positive_samples_parallel(
num_samples=config.size,
random=config.random,
seed=config.seed,
full=config.full
)
dataset.save(config.output)
if __name__ == "__main__":
# Parse command-line arguments
from src.utils.io_utils import PROJECT_ROOT
parser = argparse.ArgumentParser(description="Generate and save a dataset based on the given configuration.")
parser.add_argument("-s", "--size", type=int, default=10, help="Number of samples to generate.")
parser.add_argument("-r", "--random", type=bool, default=True, help="Whether to sample randomly.")
parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility.")
parser.add_argument("--full", action="store_true", help="Boolean flag to indicate full dataset mode.")
parser.add_argument("-o", "--output", type=str, default=os.path.join(PROJECT_ROOT, "data/dataset.pkl"), help="Output file path to save the dataset as a CSV.")
# Parse the arguments and pass to the main function
config = parser.parse_args()
main(config)