Out of Memory ?

[HadiSDev]

Hi.

Is it normal that memory should keep increasing to the point where your run out of memory? I have this example:

from functools import partial
import litdata as ld
import pandas as pd
from quran_ai.preprocessing import prepare_dataset
from transformers import (
    WhisperFeatureExtractor,
    WhisperTokenizerFast,
    WhisperProcessor,
)

model_id = "openai/whisper-large-v3-turbo"
feature_extractor = WhisperFeatureExtractor.from_pretrained(model_id)
tokenizer = WhisperTokenizerFast.from_pretrained(
    model_id, language="Arabic", task="transcribe"
)
processor = WhisperProcessor.from_pretrained(
    model_id, language="Arabic", task="transcribe"
)


if __name__ == "__main__":
    df_full = pd.read_csv("data.csv")
    df_full = df_full.drop(columns=["Unnamed: 0"])
    ld.optimize(
        fn=partial(prepare_dataset, tokenizer=tokenizer, feature_extractor=feature_extractor),                   # the function applied to each input
        inputs=df_full.to_dict(orient="records"),           # the inputs to the function (here it's a list of numbers)
        output_dir="prepared_data",             # optimized data is stored here
        num_workers=4,                      # The number of workers on the same machine
        chunk_size=1000       ,           # size of each chunk,        
    )

def prepare_dataset(batch, 
                    feature_extractor: WhisperFeatureExtractor, 
                    tokenizer: WhisperTokenizer):
    try:
    # load and resample audio data from 48 to 16kHz
        audio_file_path = batch["audio_file_path"]
        audio, sr = load_wave(audio_file_path)
        # compute log-Mel input features from input audio array 
        batch["input_features"] = feature_extractor(audio.numpy(), sampling_rate=sr).input_features[0]

        # encode target text to label ids 
        batch["labels"] = tokenizer(batch["arabic_text"]).input_ids
    except Exception as e:
        pass
    
    yield batch