I am building my own tokenizer using the byte pair encoding algorithm. I am applying this on a HuggingFace dataset. Due to memory constraints of my computer, I am first converting the dataset to a iterable_dataset which has lazy processing. Here is the part of my tokenizer function that applies to the problem.

from typing_extensions import Text
import regex as re

class Tokenizer():
  def __init__(self, vocab_size:int = 256):

    '''
      self.vocab_size:  total size of the token vocabulary
      self.num_merges:  number for merges from the base vocabulary needed
                          to reach the vocab_size.
      self.merges:      dictionary with keys: encoded integer pair of a byte pair
                                        values: associated mapping to a new encoding integer
      self.vocab:       dictionary with keys: all encoding integer numbers
                                        values: associated bytes that the integers map to.
    '''
    self.vocab_size = vocab_size
    self.num_merges = vocab_size -256
    self.merges = {} # (int,int) ->int
    self.vocab = {} # int -> bytes
    self.pat_str= r"""'(?i:[sdmt]|ll|ve|re)|[^rnp{L}p{N}]?+p{L}+|p{N}{1,3}| ?[^sp{L}p{N}]++[rn]*|s*[rn]|s+(?!S)|s+"""

    self.special_tokens = {}
    self.inverse_special_tokens = {}
    self.regular_vocab_size = None
    self.global_stats = None

  def _get_stats(self, tokens, counts = None):
    '''
    Recieves a list of tokens and returns
    a dictionary of all token pairs and the
    number of time it occurs in the string.
    '''
    count = {} if counts is None else counts
    for pair in zip(tokens, tokens[1:]):
      count[pair] = count.get(pair,0)+1 # if not found, default 0
    return count

  def _merge(self,ids, pair, idx):

    # in the list of ints (ids), replace all consecutive occurance of pair
    # with new token idx
    new_ids =[]
    i=0
    while i < len(ids):
      # if we are not at the very last position and pair matches replace it

      if i < len(ids) -1 and ids[i]==pair[0] and ids[i+1]== pair[1]:
        new_ids.append(idx)
        i+=2
      else:
        new_ids.append(ids[i])
        i+=1
    return new_ids

  def _preprocess_text(self, texts, is_batched:bool = True):
    if is_batched:
      return [re.findall(self.pat_str, text) for text in texts]
    else:
      return re.findall(self.pat_str, texts)


  def _reset_global_stats(self, reset:bool = False):
    '''
    This function is used if the global stats varaible needs to be reset. This is necessary
    when externally iterating between accumulate stats to build_vocab. After every merge,
    the stats need to be recounted.
    '''
    if reset:
      self.global_stats = None


  def accumulate_stats(self, example, index, is_batched:bool = True):
    self.global_stats = {} if self.global_stats is None else self.global_stats

    if index >0:
      token_chunks_list = example
    else:
      text_chunks_list = self._preprocess_text(example, is_batched)

      token_chunks_list = [[list(ch.encode("utf-8")) for ch in text_chunks]
                                for text_chunks in text_chunks_list ]


    for token_chunks in token_chunks_list:
      for chunk in token_chunks:
        self.global_stats = self._get_stats(chunk, self.global_stats)
  

  def save_vocab(self):
    vocab = {idx: bytes([idx]) for idx in range(256)}
    for (p0, p1), idx in self.merges.items():
      vocab[idx] = vocab[p0]+ vocab[p1]
    self.vocab = vocab # may need to convert back to dict.
    self.regular_vocab_size = len(self.vocab)

  def build_vocab(self, example, index, is_batched:bool = True):
    '''
    This has been made to be usable with huggingface datasets. Particularly the wikipedia
    dataset
    '''
    if index> 0:
      token_chunks_list = example
    else:
      text_chunks_list= self._preprocess_text(example, is_batched)

      token_chunks_list = [[list(ch.encode("utf-8")) for ch in text_chunks] for text_chunks in text_chunks_list]


    stats = self.global_stats
    # find largest value and return the associated pair that appears most frequently
    top_pair = max(stats, key = stats.get)
    idx = 256 + index
    #generatore expression to merge each chunk of tokens in list

    token_chunks = [[self._merge(chunk, top_pair, idx) for chunk in token_chunks] for token_chunks in token_chunks_list]
  
    self.merges[top_pair] = idx
    return token_chunks

And here is how I am loading the data and processing it.

from datasets import load_dataset
from datasets import *
from torch.utils.data import DataLoader
import sys
from datasets.distributed import split_dataset_by_node
sys.setrecursionlimit(2000)
from functools import partial
from datasets import Dataset


# Load a sample dataset from Hugging Face
d = load_dataset('ag_news', split='train')
d_iterable = d.to_iterable_dataset(num_shards = 50)
d_shard = split_dataset_by_node(d_iterable, world_size=50, rank=0) # this nicely gives us the very first shard only

t2 = Tokenizer(500)

def stats_build(example, index):
   # text = "".join(example["text"])
  print(index)
  t2.accumulate_stats(example["text"], index)  # Adjust this based on your dataset structure
  return example  # Replace with your tokenizer method

def build_vocab(example, index):

  example["text"] = t2.build_vocab(example["text"], index)
  return example

for i in range(500-256):

  t2._reset_global_stats(reset = True)

  j_shard = d_shard.map(stats_build, batched = True, batch_size =100, fn_kwargs = {"index": i})

  for example in j_shard:
    tokenized_text = example['text']

  d_shard=j_shard.map(build_vocab, batched = True, batch_size =100, fn_kwargs={"index":i})

  for example in d_shard:
    tokenized_text = example['text']

t2.save_vocab()

What I need this to do is run the full dataset through the accumulate_stats function. Then once it is finished processing, to run the full dataset through the build_vocab function. This needs to happen every iteration. For a sanity check, I am expecting 24 batches of the data per every loop interaction meaning I should expect 24 prints of each index before the next loop iteration. However when I include the second map function corresponding to build_vocab the print(index) line within stats_build initially prints far too many zeros before flip flopping between 0 and 1.

The map function for HuggingFace iterable datasets is normally used with multiple processing functions by having each example chain through each function before pulling the next example like so:

my_iterable_dataset = my_iterable_dataset.map(process_fn_1)
my_iterable_dataset = my_iterable_dataset.filter(filter_fn)
my_iterable_dataset = my_iterable_dataset.map(process_fn_2)

# process_fn_1, filter_fn and process_fn_2 are applied on-the-fly when iterating over the dataset
for example in my_iterable_dataset:  
    print(example)
    break

I want to try something more akin to this

my_iterable_dataset = my_iterable_dataset.map(process_fn_1)

# process_fn_1, filter_fn and process_fn_2 are applied on-the-fly when iterating over the dataset
for example in my_iterable_dataset:  
    print(example)
    
my_iterable_dataset = my_iterable_dataset.map(process_fn_2)

for example in my_iterable_dataset:  
    print(example)
    break

To do this I have tried renaming my functions and the dataset returned in several different ways as I thought that have the returned dataset having the same name from different functions might be breaking the applied-on-the-fly approach.