|
|
import os |
|
|
import logging |
|
|
import tensorflow as tf |
|
|
import pandas as pd |
|
|
import pyarrow.parquet as pq |
|
|
from sklearn.preprocessing import MinMaxScaler |
|
|
|
|
|
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
def create_data_loader(parquet_paths: list, scaler: MinMaxScaler, seq_length: int, batch_size: int) -> tf.data.Dataset: |
|
|
"""Create a tf.data.Dataset from Parquet files for LSTM training or evaluation. |
|
|
|
|
|
Args: |
|
|
parquet_paths (list): List of paths to Parquet files. |
|
|
scaler (MinMaxScaler): Scaler fitted on the data. |
|
|
seq_length (int): Length of input sequences. |
|
|
batch_size (int): Batch size for the dataset. |
|
|
|
|
|
Returns: |
|
|
tf.data.Dataset: Dataset yielding (sequence, target) pairs with shapes (batch_size, seq_length, 1) and (batch_size, 1). |
|
|
|
|
|
Raises: |
|
|
ValueError: If inputs are invalid or no valid data is found. |
|
|
""" |
|
|
if not parquet_paths: |
|
|
logger.error("No parquet paths provided") |
|
|
raise ValueError("parquet_paths cannot be empty") |
|
|
if not isinstance(scaler, MinMaxScaler): |
|
|
logger.error("Invalid scaler provided") |
|
|
raise ValueError("scaler must be an instance of MinMaxScaler") |
|
|
if not isinstance(seq_length, int) or seq_length <= 0: |
|
|
logger.error(f"Invalid seq_length: {seq_length}") |
|
|
raise ValueError("seq_length must be a positive integer") |
|
|
if not isinstance(batch_size, int) or batch_size <= 0: |
|
|
logger.error(f"Invalid batch_size: {batch_size}") |
|
|
raise ValueError("batch_size must be a positive integer") |
|
|
|
|
|
total_sequences = 0 |
|
|
def _scaled_generator(): |
|
|
nonlocal total_sequences |
|
|
for path in parquet_paths: |
|
|
if not os.path.exists(path): |
|
|
logger.warning(f"Parquet file not found, skipping: {path}") |
|
|
continue |
|
|
try: |
|
|
file_size = os.path.getsize(path) / (1024 * 1024) |
|
|
if file_size < 100: |
|
|
df = pd.read_parquet(path, columns=['Close']) |
|
|
logger.debug(f"Loaded {path} into memory, size: {file_size:.2f} MB") |
|
|
if 'Close' not in df.columns or df['Close'].isna().any(): |
|
|
logger.warning(f"Invalid or missing 'Close' column in {path}") |
|
|
continue |
|
|
prices = df['Close'].astype('float32').values.reshape(-1, 1) |
|
|
if prices.size <= seq_length: |
|
|
logger.warning(f"File {path} has {prices.size} rows, insufficient for seq_length {seq_length}") |
|
|
continue |
|
|
scaled = scaler.transform(prices) |
|
|
for j in range(len(scaled) - seq_length): |
|
|
total_sequences += 1 |
|
|
yield scaled[j:j + seq_length], scaled[j + seq_length] |
|
|
else: |
|
|
parquet_file = pq.ParquetFile(path) |
|
|
for batch in parquet_file.iter_batches(batch_size=10_000, columns=['Close']): |
|
|
chunk = batch.to_pandas() |
|
|
if 'Close' not in chunk.columns or chunk['Close'].isna().any(): |
|
|
logger.warning(f"Invalid or missing 'Close' column in {path}") |
|
|
continue |
|
|
prices = chunk['Close'].astype('float32').values.reshape(-1, 1) |
|
|
scaled = scaler.transform(prices) |
|
|
logger.debug(f"Processing batch from {path}, scaled shape: {scaled.shape}") |
|
|
for j in range(len(scaled) - seq_length): |
|
|
total_sequences += 1 |
|
|
yield scaled[j:j + seq_length], scaled[j + seq_length] |
|
|
except Exception as e: |
|
|
logger.error(f"Error processing parquet file {path}: {e}") |
|
|
continue |
|
|
|
|
|
if total_sequences == 0: |
|
|
logger.error("No valid sequences generated from any Parquet file") |
|
|
raise ValueError("No valid sequences generated from any Parquet file") |
|
|
|
|
|
dataset = tf.data.Dataset.from_generator( |
|
|
_scaled_generator, |
|
|
output_types=(tf.float32, tf.float32), |
|
|
output_shapes=((seq_length, 1), (1,)) |
|
|
).batch(batch_size).prefetch(tf.data.AUTOTUNE) |
|
|
|
|
|
logger.info(f"Created data loader with seq_length={seq_length}, batch_size={batch_size}, total_sequences={total_sequences}") |
|
|
return dataset |
