homework7/count_vectorizer.py - hw7 - azatgimaev/python_course

homework7/count_vectorizer.py
106 строк3 KB
Azat Gimaev
Допуск
19 ноя 2024, 02:49
19 ноя 2024, 02:49382b017
def split_string_into_words(text: str, lowercase: bool) -> list:
    '''
    This function splits the text into a list of words.
    The result does not include characters from string.punctuation

    Parameters
    ----------------
    :param text: str
        The text to be divided into words
    :param lowercase: bool
        -If True:
            The text will be changed to lowercase.
            The words will be in lowercase
        -If False:
            the text will not be changed

    Return
    -----------------
    :return: list
        The list of words in the text
    '''
    if lowercase:
        text = text.lower()
    for i in '!"#$%&()*+,-./:;<=>?@[]^_`{|}~':
        text = text.replace(i, '')
    return text.split()


class CountVectorizer:
    def __init__(
        self,
        lowercase=True
    ):
        self.vocabulary = None
        self.lowercase = lowercase

    def _set_vocabulary(self, raw_texts: list):
        '''
        The function finds all words from raw_texts
        and assigns this list of words to vocabulary

        Parameters
        -----------------
        :param raw_texts: list
            An iterable which generates either str
        '''
        words_in_raw_texts = []

        if not isinstance(raw_texts, list):
            raise ValueError('The function takes a list of strings as input')

        for text in raw_texts:
            if not isinstance(text, str):
                raise ValueError('The \'' + str(text) + '\' is not a string')
            for word in split_string_into_words(text, self.lowercase):
                if word not in words_in_raw_texts:
                    words_in_raw_texts.append(word)

        self.vocabulary = words_in_raw_texts

    def fit_transform(self, raw_texts: list) -> list:
        '''
        Learn the vocabulary dictionary and return document-term matrix.

        Parameters
        -----------------
        :param raw_texts: iterable of str
            An iterable which generates either str

        Return
        ----------------
        :return: list
            Document-term matrix.
        '''
        self._set_vocabulary(raw_texts)

        transform_array = []
        for text in raw_texts:
            transform_text_array = []
            for word in self.vocabulary:
                text_words = split_string_into_words(text, self.lowercase)
                count = text_words.count(word)
                transform_text_array.append(count)
            transform_array.append(transform_text_array)
        return transform_array

    def get_feature_names(self) -> list:
        '''
        Returns a list of words

        Return
        ---------------
        :return: list
            vocabulary
        '''
        return self.vocabulary


if __name__ == "__main__":
    a = CountVectorizer()
    corpus = [
        'Crock Pot Pasta Never boil pasta again',
        'Pasta Pomodoro Fresh ingredients Parmesan to taste'
    ]
    print(a.fit_transform(corpus))
    print(a.get_feature_names())