from datetime import datetime from tqdm import tqdm from datasets import load_dataset from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor from items import Item CHUNK_SIZE = 1000 MIN_PRICE = 0.5 MAX_PRICE = 999.49 class ItemLoader: def __init__(self, name): self.name = name self.dataset = None def from_datapoint(self, datapoint): """ Try to create an Item from this datapoint Return the Item if successful, or None if it shouldn't be included """ try: price_str = datapoint['price'] if price_str: price = float(price_str) if MIN_PRICE <= price <= MAX_PRICE: item = Item(datapoint, price) return item if item.include else None except ValueError: return None def from_chunk(self, chunk): """ Create a list of Items from this chunk of elements from the Dataset """ batch = [] for datapoint in chunk: result = self.from_datapoint(datapoint) if result: batch.append(result) return batch def chunk_generator(self): """ Iterate over the Dataset, yielding chunks of datapoints at a time """ size = len(self.dataset) for i in range(0, size, CHUNK_SIZE): yield self.dataset.select(range(i, min(i + CHUNK_SIZE, size))) def load_in_parallel(self, workers): """ Use concurrent.futures to farm out the work to process chunks of datapoints - This speeds up processing significantly, but will tie up your computer while it's doing so! """ results = [] chunk_count = (len(self.dataset) // CHUNK_SIZE) + 1 with ProcessPoolExecutor(max_workers=workers) as pool: for batch in tqdm(pool.map(self.from_chunk, self.chunk_generator()), total=chunk_count): results.extend(batch) for result in results: result.category = self.name return results def load(self, workers=8): """ Load in this dataset; the workers parameter specifies how many processes should work on loading and scrubbing the data """ start = datetime.now() print(f"Loading dataset {self.name}", flush=True) self.dataset = load_dataset("McAuley-Lab/Amazon-Reviews-2023", f"raw_meta_{self.name}", split="full", trust_remote_code=True) results = self.load_in_parallel(workers) finish = datetime.now() print(f"Completed {self.name} with {len(results):,} datapoints in {(finish-start).total_seconds()/60:.1f} mins", flush=True) return results