tokenize_data.py

from dotenv import load_dotenv
load_dotenv()

from transformers import AutoTokenizer
from datasets import load_dataset, load_from_disk, concatenate_datasets
from concurrent.futures import ThreadPoolExecutor, as_completed
from tqdm import tqdm
import time
import os

# Load a pre-trained tokenizer
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B")
tokenizer.pad_token = tokenizer.eos_token

# Define all dataset loading tasks
def load_wikitext_103_train():
    ds = load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1", split="train")
    ds = ds.filter(lambda x: len(x["text"].strip()) > 0)
    return ds.select_columns(["text"])

def load_wikitext_103_val_test():
    parts = []
    for split in ["validation", "test"]:
        ds = load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1", split=split)
        ds = ds.filter(lambda x: len(x["text"].strip()) > 0)
        parts.append(ds.select_columns(["text"]))
    return concatenate_datasets(parts)

def load_wikitext_2():
    parts = []
    for split in ["train", "validation", "test"]:
        ds = load_dataset("Salesforce/wikitext", "wikitext-2-raw-v1", split=split)
        ds = ds.filter(lambda x: len(x["text"].strip()) > 0)
        parts.append(ds.select_columns(["text"]))
    return concatenate_datasets(parts)

def load_english_dict():
    ds = load_dataset("npvinHnivqn/EnglishDictionary", split="train")
    ds = ds.map(lambda x: {"text": f"{x['word']}: {x['definition']}"})
    return ds.select_columns(["text"])

def load_wordnet():
    ds = load_dataset("marksverdhei/wordnet-definitions-en-2021", split="train")
    ds = ds.map(lambda x: {"text": f"{x['Word']}: {x['Definition']}. Example: {x['Example']}"})
    return ds.select_columns(["text"])

def load_ag_news():
    ds = load_dataset("fancyzhx/ag_news", split="train")
    return ds.select_columns(["text"])

def load_imdb():
    ds = load_dataset("stanfordnlp/imdb", split="train")
    return ds.select_columns(["text"])

def load_rotten_tomatoes():
    ds = load_dataset("cornell-movie-review-data/rotten_tomatoes", split="train")
    return ds.select_columns(["text"])

def load_cnn_dailymail():
    ds = load_dataset("abisee/cnn_dailymail", "3.0.0", split="train")
    ds = ds.rename_column("article", "text")
    return ds.select_columns(["text"])

def load_yelp():
    ds = load_dataset("Yelp/yelp_review_full", split="train")
    return ds.select_columns(["text"])

def load_urban_dictionary():
    ds = load_dataset("daspartho/urban_dictionary", split="train")
    ds = ds.map(lambda x: {"text": f"{x['word']}: {x['definition']}. Example: {x['example']}"})
    return ds.select_columns(["text"])

def load_slang():
    ds = load_dataset("LM-Lexicon/Slang", split="train")
    ds = ds.map(lambda x: {"text": f"{x['term']}: {x['definition']}. Example: {x['context']}"})
    return ds.select_columns(["text"])

def load_genz_slang():
    ds = load_dataset("MLBtrio/genz-slang-dataset", split="train")
    ds = ds.map(lambda x: {"text": f"{x['Slang']}: {x['Description']}. Example: {x['Example']}"})
    return ds.select_columns(["text"])

# All tasks with labels
tasks = [
    ("WikiText-103 (train)", load_wikitext_103_train),
    ("WikiText-103 (val+test)", load_wikitext_103_val_test),
    ("WikiText-2", load_wikitext_2),
    ("English Dictionary", load_english_dict),
    ("WordNet Definitions", load_wordnet),
    ("AG News", load_ag_news),
    ("IMDB", load_imdb),
    ("Rotten Tomatoes", load_rotten_tomatoes),
    ("CNN/DailyMail", load_cnn_dailymail),
    ("Yelp Reviews", load_yelp),
    ("Urban Dictionary", load_urban_dictionary),
    ("LM-Lexicon Slang", load_slang),
    ("Gen Z Slang", load_genz_slang),
]

# Directory where datasets are cached locally
data_dir = os.path.join(os.path.dirname(__file__), "data")
os.makedirs(data_dir, exist_ok=True)

def safe_name_for(label):
    return label.lower().replace(" ", "_").replace("/", "_").replace("(", "").replace(")", "").replace("+", "_")

print(f"Loading {len(tasks)} datasets...\n")

# Create one tqdm bar per dataset, each on its own line
bars = []
for i, (label, _) in enumerate(tasks):
    bar = tqdm(total=1, desc=f"  {i+1:>2}/{len(tasks)} {label:<25}", position=i, leave=True,
               bar_format="{desc} {bar} {postfix}")
    bar.set_postfix_str("pending...")
    bars.append(bar)

results = {}

def run_task(index, label, fn):
    # Check if dataset already exists on disk
    save_path = os.path.join(data_dir, safe_name_for(label))
    if os.path.isdir(save_path):
        try:
            ds = load_from_disk(save_path)
            bars[index].update(1)
            bars[index].set_postfix_str(f"✓ {len(ds)} examples (cached)")
            return label, ds
        except Exception:
            pass  # cache corrupt, re-download

    # Download from HuggingFace
    bars[index].set_postfix_str("downloading...")
    max_retries = 5
    for attempt in range(max_retries):
        try:
            ds = fn()
            # Save to disk for next time
            ds.save_to_disk(save_path)
            bars[index].update(1)
            bars[index].set_postfix_str(f"✓ {len(ds)} examples")
            return label, ds
        except Exception as e:
            if "429" in str(e) or "rate limit" in str(e).lower():
                wait_time = 60 * (attempt + 1)
                bars[index].set_postfix_str(f"rate limited, retry in {wait_time}s...")
                time.sleep(wait_time)
            else:
                bars[index].set_postfix_str(f"error, retry {attempt+1}/{max_retries}...")
                time.sleep(10 * (attempt + 1))
            if attempt == max_retries - 1:
                bars[index].set_postfix_str(f"✗ FAILED: {e}")
                raise

with ThreadPoolExecutor(max_workers=3) as executor:
    futures = {executor.submit(run_task, i, label, fn): label for i, (label, fn) in enumerate(tasks)}
    for future in as_completed(futures):
        label, ds = future.result()
        results[label] = ds

# Close all bars and move cursor below them
for bar in bars:
    bar.close()
print(f"\nAll {len(tasks)} datasets loaded!")

# Collect in original order
datasets_list = [results[label] for label, _ in tasks]

# Combine all datasets
print("\nCombining datasets...")
combined_dataset = concatenate_datasets(datasets_list)
print(f"Total examples: {len(combined_dataset)}")

# Tokenize dataset
def tokenize_fn(examples):
    return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=128)

print("Tokenizing...")
tokenized_dataset = combined_dataset.map(tokenize_fn, batched=True)
print("Done!")