Draw back : lets assume in some file one line context and  5 line spaces are there (may or may not)
2) all tensor should same size (if data is not there some padding tragies)
3) returns input_ids, attention_mask, labels

4) data collection and loading stragies are very  old , not novelity missing



import os
from typing import List, Dict
from transformers import PreTrainedTokenizer, AutoTokenizer


def read_txt_files(directory: str) -> List[str]:
    """
    Reads all .txt files from the given directory and returns a list of their contents.

    Args:
        directory (str): The path to the directory containing .txt files.

    Returns:
        List[str]: A list of strings, each string is the content of a .txt file.
    """
    contents = []
    try:
        for filename in os.listdir(directory):
            if filename.endswith('.txt'):
                file_path = os.path.join(directory, filename)
                with open(file_path, 'r', encoding='utf-8') as file:
                    content = file.read()
                    # Clean up the content by stripping whitespace and handling empty lines
                    content = '\n'.join(line.strip() for line in content.splitlines() if line.strip())
                    contents.append(content)
    except Exception as e:
        print(f"An error occurred while reading files: {e}")
    return contents


def split_text(
    text: str,
    tokenizer: PreTrainedTokenizer,
    max_length: int
) -> List[str]:
    """
    Splits a text into chunks that fit within the max_length when tokenized.

    Args:
        text (str): The text to split.
        tokenizer (PreTrainedTokenizer): The tokenizer to use.
        max_length (int): The maximum length of the tokenized output.

    Returns:
        List[str]: A list of text chunks.
    """
    words = text.strip().split()
    chunks = []
    current_chunk = []
    current_length = 0
    for word in words:
        # Encode the word to get its token length
        tokenized_word = tokenizer.encode(word, add_special_tokens=False)
        word_length = len(tokenized_word)
        # Check if adding the word exceeds the max_length
        if current_length + word_length > max_length:
            if current_chunk:
                chunks.append(' '.join(current_chunk))
            current_chunk = [word]
            current_length = word_length
        else:
            current_chunk.append(word)
            current_length += word_length
    if current_chunk:
        chunks.append(' '.join(current_chunk))
    return chunks


def tokenize_texts(
    texts: List[str],
    tokenizer: PreTrainedTokenizer,
    max_length: int
) -> Dict[str, List[List[int]]]:
    """
    Tokenizes a list of texts using the given tokenizer, ensuring zero data loss by splitting long texts.
    Adds BOS and EOS tokens manually.

    Args:
        texts (List[str]): The texts to tokenize.
        tokenizer (PreTrainedTokenizer): The tokenizer to use.
        max_length (int): The maximum length of the tokenized output.

    Returns:
        Dict[str, List[List[int]]]: A dictionary containing 'input_ids' and 'attention_mask'.
    """
    all_input_ids = []
    all_attention_masks = []
    for text in texts:
        # Skip empty texts
        if not text:
            continue
        chunks = split_text(text, tokenizer, max_length - 2)  # Reserve space for BOS and EOS tokens
        for chunk in chunks:
            # Tokenize the chunk without adding any special tokens
            tokens = tokenizer.encode(
                chunk,
                add_special_tokens=False,
                max_length=max_length - 2,
                truncation=True
            )
            # Manually add BOS and EOS tokens
            tokens = [tokenizer.bos_token_id] + tokens + [tokenizer.eos_token_id]
            # Create attention mask
            attention_mask = [1] * len(tokens)
            # Pad if necessary
            if len(tokens) < max_length:
                padding_length = max_length - len(tokens)
                tokens += [tokenizer.pad_token_id] * padding_length
                attention_mask += [0] * padding_length
            else:
                tokens = tokens[:max_length]
                attention_mask = attention_mask[:max_length]
            all_input_ids.append(tokens)
            all_attention_masks.append(attention_mask)
    return {'input_ids': all_input_ids, 'attention_mask': all_attention_masks}


def main():
    directory = '/scratch/hemanth/Hemanth/localdataset'  # Replace with your directory path
    max_length = 512  # Define your desired max sequence length

    # Initialize the tokenizer
    tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')

    # Add special tokens if they are not already present
    if not tokenizer.special_tokens_map.get('bos_token'):
        tokenizer.add_special_tokens({'bos_token': '[BOS]'})
    if not tokenizer.special_tokens_map.get('eos_token'):
        tokenizer.add_special_tokens({'eos_token': '[EOS]'})
    if not tokenizer.special_tokens_map.get('pad_token'):
        tokenizer.add_special_tokens({'pad_token': '[PAD]'})

    # Adjust the tokenizer to recognize the new special tokens
    tokenizer.bos_token = '[BOS]'
    tokenizer.eos_token = '[EOS]'
    tokenizer.pad_token = '[PAD]'

    # Update tokenizer's vocabulary to include the new tokens
    tokenizer.save_pretrained('./tokenizer_with_special_tokens')
    tokenizer = AutoTokenizer.from_pretrained('./tokenizer_with_special_tokens')

    # Read the texts from files
    texts = read_txt_files(directory)

    # Tokenize the texts
    tokenized_data = tokenize_texts(texts, tokenizer, max_length)
    print(len(tokenized_data['input_ids']))

    # You can now use tokenized_data['input_ids'] and tokenized_data['attention_mask'] for your model
    print('Tokenization complete.')
    # For example, to see the first tokenized input:
    print(tokenized_data['input_ids'][0])
    print(tokenized_data['attention_mask'][0])


if __name__ == '__main__':
    main()