One of my questions was how to set batch_size param in the DistilBertModel.
I implemented a naive code to split my data into batches (please see below). But I am sure there is a better way by using some ready to go API. Any idea?
import torch

def batch_embeddings(tokenizer, texts, batch_size=256, max_length=MAX_LENGTH):
    all_cls_tokens = []
    
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i+batch_size]
        tokenizer_output = tokenizer.batch_encode_plus(batch,
                                             max_length=max_length,
                                             padding='longest',
                                             truncation=True,
                                             return_attention_mask=True,
                                             return_token_type_ids=False,
                                             return_tensors="pt"
                                             )
        cls_token = model(**tokenizer_output)[0][:, 0, :]

        all_cls_tokens = all_cls_tokens + cls_token.tolist()
        # here I would save all_cls_tokens into disk time to time and clean it up. It would avoid getting the same ram error as previous code

    return all_cls_tokens


all_cls_tokens = batch_embeddings(tokenizer, pd_df['text'].tolist(), 1000)


Also I see several issues here since I am doing several conversions:
Initially my data is in a pyspark dataframe (I don’t show it here to simplify the question). I converted it to pandas dataframe, then to a list. Inside method batch_embeddings, I sliced this list (variable called batch) and passed it into method tokenizer.batch_encode_plus, which return a tensor.torch (variable called cls_token). I converted it to a list and append it to the previous result. This is another reason I am sure there is another solution.
What would be a better approach to deal with a lot of data using DistilBertModel?
Please help!
Thanks!