Split Text into Tokenized Chunks

split_sentence2maxtok.py

#!/usr/bin/env python3
#python3 -m pip install nltk tiktoken
#chmod +x split_sentence2maxtok.py
#./split_sentence2maxtok.py "This is a test sentence. This is another test sentence. This is a long test sentence." 6

import nltk
from nltk.tokenize import sent_tokenize
from nltk.data import find
import tiktoken
import sys

def get_nltk_punkt():
    try:
        # Check if the 'punkt' package has already been downloaded
        find('tokenizers/punkt')
    except LookupError:
        try:
            # Attempt to download the 'punkt' package
            nltk.download('punkt')
        except Exception as e:
            # Handle the download error
            print(f"Error downloading the 'punkt' package: {e}")
            sys.exit(1)
            
def split_sentences(text, max_length=512):
    get_nltk_punkt()
    sentences = sent_tokenize(text)
    enc = tiktoken.get_encoding("gpt2")
    short_sentences = []
    for sentence in sentences:
        tokens = enc.encode(sentence)
        while len(tokens) > max_length:
            current_sentence = tokens[:max_length]
            # Ensure we don't split in the middle of a word or punctuation
            if len(current_sentence) < len(tokens):
                split_point = max_length
                while split_point > 0 and not enc.decode([tokens[split_point]]).isspace():
                    split_point -= 1
                if split_point == 0:
                    split_point = max_length
                current_sentence = tokens[:split_point]
            short_sentences.append(enc.decode(current_sentence).strip())
            tokens = tokens[len(current_sentence):]
        short_sentences.append(enc.decode(tokens).strip())
    
    return short_sentences

def main():
    if len(sys.argv) < 2:
        print("Usage: python script.py <text> [max_length]")
        sys.exit(1)
    
    text = sys.argv[1]
    max_length = int(sys.argv[2]) if len(sys.argv) > 2 else 512
    
    sentences = split_sentences(text, max_length)
    for sentence in sentences:
        print(sentence)

if __name__ == "__main__":
    main()