<jupyter_start><jupyter_text>Since mapping dataset takes a long (~15min) time, it is done separately and then saved, so that the training can be done by directly loading the already mapped dataset.<jupyter_code># Following SimCSE procedure.
import pandas as pd
from datasets import Dataset,load_dataset
from transformers import BertTokenizer, BertModel, AdamW, get_linear_schedule_with_warmup
from torch.utils.data import DataLoader
import torch
import torch.nn.functional as F
import torch.optim as optim
import torch.optim.lr_scheduler as lr_scheduler
from tqdm import tqdm
import torch.nn as nn
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')<jupyter_output><empty_output><jupyter_text>Download wiki1m_for_simCSE.txt:https://huggingface.co/datasets/princeton-nlp/datasets-for-simcse/resolve/main/wiki1m_for_simcse.txt<jupyter_code># use pandas to read simCSE-wiki.txt
wiki_text_file = 'your-path-to/simCSE-wiki.txt'
wiki = pd.read_csv(wiki_text_file,sep = '\t',header = None)
wiki.columns = ['text']
# use Dataset.from_pandas to convert to dataset
wiki_dataset = Dataset.from_pandas(wiki,split= "train")
wiki_dataset
def prepare_features(examples):
    
    total = len(examples['text'])
    # total = batch_size
    
    # Avoid "None" fields 
    for idx in range(total):
        if examples['text'][idx] is None:
            examples['text'][idx] = " "
        if examples['text'][idx] is None:
            examples['text'][idx] = " "

    sentences = examples['text'] + examples['text']

    # set max_length here:
    sent_features = tokenizer(sentences, max_length=32, truncation=True, padding="max_length")

    features = {}
    for key in sent_features:
        features[key] = [[sent_features[key][i], sent_features[key][i+total]] for i in range(total)]
    
    return features
train_dataset = wiki_dataset.map(prepare_features,batched=True, remove_columns=['text'], batch_size=4000)
train_dataset.set_format(type='torch', columns=['input_ids', 'token_type_ids', 'attention_mask'])
# save to disk for reuse
train_dataset.save_to_disk("wiki_for_sts_32")<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11084793936522976, 'fraction_numerical': 0.005684509711037423, 'mean_word_length': 5.267062314540059, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '282493', 'license': 'MIT', 'n_tokens_mistral': 695, 'n_tokens_neox': 679, 'n_words': 161, 'path': 'data/wiki1m.ipynb', 'repo_name': 'AdamRain/mocose', 'size': '2111'}