# custom preprocessing class with bigram generation
class NgramPreprocessing:
    def __init__(self, ngram_range=(1, 1), 
                       vocab_size=10000, 
                       stopwords='English'): 

        self.ngram_range = ngram_range
        self.preprocessing = Preprocessing(vocab_size=vocab_size, 
                                           stopwords=stopwords)

        # use a custom analyzer to join bigrams with "_"
        self.vectorizer = CountVectorizer(ngram_range=self.ngram_range, 
                                          max_features=vocab_size, 
                                          analyzer=self._custom_analyzer)

        # custom analyzer function to join bigrams with underscores
    def _custom_analyzer(self, doc):
        # tokenize the document and create bigrams
        tokens = CountVectorizer(ngram_range=self.ngram_range).build_analyzer()(doc)

        # replace spaces in bigrams with "_"
        return [token.replace(" ", "_") for token in tokens]

    def preprocess(self, 
                   docs, 
                   pretrained_WE=False):

        parsed_docs = self.preprocessing.preprocess(docs, 
                      pretrained_WE=pretrained_WE)["train_texts"]
        train_bow = self.vectorizer.fit_transform(parsed_docs).toarray()
        rst = {
            "train_bow": train_bow,
            "train_texts": parsed_docs,
            "vocab": self.vectorizer.get_feature_names_out()
        }
        return rst

# initialize preprocessing with bigrams
ngram_preprocessing = NgramPreprocessing(ngram_range=(2, 2))