From the uDemy course on LLM engineering.
https://www.udemy.com/course/llm-engineering-master-ai-and-large-language-models
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
81 lines
2.7 KiB
81 lines
2.7 KiB
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 |
|
|
|
|
|
|
|
|