anveshplus/BPE-Tokenizer

1import pandas as pd2import re3 4# Data section  start--> 5# Load the CSV files6file_paths = [7    '/Users/anvesh/codebase/llm/data/telugu_books/telugu_books.csv',8    '/Users/anvesh/codebase/llm/data/telugu_news/1_telugu_news.csv',9    '/Users/anvesh/codebase/llm/data/telugu_news/2_telugu_news.csv'10]11 12# Combine data from all files13telugu_texts = []14for file_path in file_paths:15    df = pd.read_csv(file_path)16    if 'text' in df.columns:17        telugu_texts.append(' '.join(df['text'].astype(str).tolist()))18    elif 'body' in df.columns:19        telugu_texts.append(' '.join(df['body'].astype(str).tolist()))20 21# Concatenate all texts and remove all English, numerical values, and quotes22telugu_text = ' '.join(telugu_texts)23telugu_text = re.sub(r'[A-Za-z0-9\'"]', '', telugu_text)  # Remove English letters, numbers, and quotes24telugu_text = re.sub(r'[\r\n\xa0]', '', telugu_text)  # Remove line breaks and non-breaking spaces25 26print('telugu_text befores utf-8 encoding:', telugu_text[:100])27 28vocabulary_size = len(set(telugu_text.split()))29print('Original text size:', len(telugu_text))30print('Vocabulary size of telugu_text:', vocabulary_size)31 32unique_characters = set(telugu_text)33unique_count = len(unique_characters)34print('Original text size:', len(telugu_text))35print('Unique character count in telugu_text:', unique_count)36 37# Data section  end--> 38 39# utf-8 encoding section start -->40import encode_parallel_telugu as encode_parallel41import time42 43tokens = encode_parallel.load_telugu_texts()44# Start the timer45start_time = time.time()46# Encode the tokens in parallel and get concatenated results47encoded_tokens = encode_parallel.encode_tokens_parallel(tokens, chunk_size=1_000_000, max_workers=10)48print('encoded_tokens:', encoded_tokens[:100])49print(len(encoded_tokens))50# End the timer51end_time = time.time()52print(f"Time taken to encode and process tokens in parallel: {end_time - start_time:.4f} seconds")53 54print('length of encoded_text:', len(encoded_tokens))55print('unique characters in encoded_text:', set(encoded_tokens))56print('unique characters in encoded_text:', len(set(encoded_tokens)))57# utf-8 encoding section end -->58 59# BPE section start -->60#### **BPE implementation**61 62tokens = encoded_tokens63 64def get_stats(ids):65    counts = {}66    for pair in zip(ids, ids[1:]):67        counts[pair] = counts.get(pair, 0) + 168    return counts69 70def merge(ids, pair, idx):71    new_ids = []72    i = 073    while i < len(ids):74        if i < len(ids) - 1 and ids[i] == pair[0] and ids[i+1] == pair[1]:75            new_ids.append(idx)76            i += 277        else:78            new_ids.append(ids[i])79            i += 180    return new_ids81 82# ---83vocab_size = 500 # the desired final vocabulary size84num_merges = vocab_size - 256 ## our unique tokens are 194, for our sample text.85ids = list(tokens) # copy so we don't destroy the original list86 87merges = {} # (int, int) -> int88from tqdm import tqdm  # Import tqdm for progress bar89 90for i in tqdm(range(num_merges), desc="Merging tokens"):91    stats = get_stats(ids)92    pair = max(stats, key=stats.get)93    idx = 256 + i94    # print(f"merging {pair} into a new token {idx}")95    ids = merge(ids, pair, idx)96    merges[pair] = idx # merge has a pair of tokens and the new token index97    98print("tokens length:", len(tokens))99print("ids length:", len(ids))100print(f"compression ratio: {len(tokens) / len(ids):.2f}X")101print(f"token size: {len(set(tokens))}")102    103# print(ids)104# BPE section end -->105 106# Building the vocabulary section start -->107telugu_unicode_chars = [chr(i) for i in range(0x0C00, 0x0C7F)]  # Telugu Unicode range108 109# Add these characters to the vocabulary110import json111vocab = {token: idx for token, idx in merges.items()}112# Add unique Telugu characters to the vocabulary113for idx, char in enumerate([chr(i).encode('utf-8') for i in range(0x0C00, 0x0C7F)]):114    if idx < 256:  # Ensure we only add up to 256 characters115        vocab[char] = idx  # Map the character to its index116 117vocab[b' '] = 255118vocab[b'.'] = 254119# Save merges and vocab to a file120# with open('merges_vocab.json', 'w') as f:121#     json.dump({'merges': merges, 'vocab': vocab}, f)122 123# saving the merges and vocab to a file124with open('merges_vocab.json', 'w') as f:125    json.dump({'merges': {str(k): v for k, v in merges.items()}, 'vocab': {str(k): v for k, v in vocab.items()}}, f)126    127# Building the vocabulary section end -->128 129 130# Reading the merges and vocab from a file section start -->131import json132from collections import defaultdict133 134# Read the merges and vocab data from the JSON file135with open('merges_vocab.json', 'r') as f:136    data = json.load(f)137 138# Create a defaultdict to store the data in a distributed manner139distributed_data = defaultdict(list)140 141# Distribute the merges and vocab data142# for key, value in data['merges'].items():143#     distributed_data['merges'].append({key: value})144 145for key, value in data['vocab'].items():146    distributed_data['vocab'].append({key: value})147 148# Optionally, print the distributed data for verification149print(distributed_data)150distributed_data['vocab']151# Convert the list of dictionaries to a single dictionary152formatted_vocab = {}153for item in distributed_data['vocab']:154    for k, v in item.items():155        if ',' not in k:156            formatted_vocab[(eval(k),)] = v157        else:158            formatted_vocab[eval(k)] = v159print(formatted_vocab[:50])160# inverting the vocab161inverted_vocab = {v: k for k, v in formatted_vocab.items()}162inverted_vocab163 164# Reading the merges and vocab from a file section end -->165 166# Expanding the vocab section start -->167def convert_to_bytes(value):168    if isinstance(value, bytes):169        return value170    elif value in inverted_vocab:171        return process_tuple(inverted_vocab[value])172    else:173        print(f'value not found in inverted_vocab: {value}')174        return None175 176def process_tuple(value_tuple):177    # print(f'value_tuple: {value_tuple}')178    # for vi in value_tuple:179    #     print(f'v: {vi}')180    converted_values = []181    for v in value_tuple:182        result = convert_to_bytes(v)183        if isinstance(result, tuple):184            converted_values.extend(result)185        else:186            converted_values.append(result)187    return tuple(converted_values)188 189decoder_map = {k: process_tuple(v) for k, v in inverted_vocab.items()}190 191 192 193 194 195