2-2 N-그램

import re

def n_gram(text, n):
    #uni_gram
    text_low = text.lower()
    text_low = re.sub('(\W+)', ' ', text_low)
    words = text_low.split(' ')

    if '' in words:
        words.remove('')
    
    #bi_gram
    if n == 2:
        bi_words = []
        for i in range(len(words) - 1):
            bi_words.append(' '.join(words[i:i+2]))
        return bi_words

    #n_gram
    n_words = []
    for i in range(len(words) - n + 1):
        n_words.append(' '.join(words[i:i+n]))
    return n_words


#예시
text = "Yes, I can do it!"

>>> n_gram(text, 1) #uni_gram
['yes', 'i', 'can', 'do', 'it']

>>> n_gram(text, 2) #bi_gram
bigram = ['yes i', 'i can', 'can do', 'do it']

>>> n_gram(text, 3) #tri_gram
trigram = ['yes i can', 'i can do', 'can do it']