classes.py - functional - Tailord/omd

classes.py
72 строки2 KB
Daniil Polyakov
implement TfidfVectorizer
23 окт 2024, 22:55
23 окт 2024, 22:558b3e570
from typing import Self
from itertools import chain
from math import log

class CountVectorizer:
    def __init__(self):
        self.features = None

    def fit(self, texts: list[str]) -> Self:
        texts = [set(text.lower().split()) for text in texts]
        self.features = dict.fromkeys(chain.from_iterable(texts), 0)
        return self

    def transform(self, texts: list[str]) -> list[list[int]]:
        count_matrix = []
        for text in texts:
            count_matrix.append(self.features.copy())
            for word in text.split():
                count_matrix[-1][word.lower()] += 1
        return [list(count_dict.values()) for count_dict in count_matrix]

    def fit_transform(self, texts: list[str]) -> list[list[int]]:
        return self.fit(texts).transform(texts)

    def get_feature_names(self) -> list[str]:
        return list(self.features)

def tf_transform(count_matrix: list[list[int]]) -> list[list[float]]:
    tf = []
    for row in count_matrix:
        s = sum(row)
        tf.append([term / s for term in row])
    return tf

def idf_transform(count_matrix: list[list[int]]) -> list[float]:
    n, m = len(count_matrix), len(count_matrix[0])
    if_exists_matrix = [[bool(count_matrix[i][j]) for i in range(n)] for j in range(m)]
    idf = [log((len(term) + 1) / (sum(term) + 1)) + 1 for term in if_exists_matrix]
    return idf

class TfidfTransformer:
    def __init__(self):
        self.idf = None

    def fit(self, count_matrix: list[list[int]]) -> Self:
        self.idf = idf_transform(count_matrix)
        return self

    def transform(self, count_matrix: list[list[int]]) -> list[list[float]]:
        tf = tf_transform(count_matrix)
        tfidf = []
        for tf_row in tf:
            tfidf.append([tf * idf for tf, idf in zip(tf_row, self.idf)])
        return tfidf

    def fit_transform(self, count_matrix: list[list[int]]) -> list[list[float]]:
        return self.fit(count_matrix).transform(count_matrix)

class TfidfVectorizer(CountVectorizer):
    def __init__(self):
        super().__init__()
        self.transformer = TfidfTransformer()

    def fit(self, texts: list[str]) -> Self:
        super().fit(texts)
        count_matrix = super().transform(texts)
        self.transformer.fit(count_matrix)
        return self

    def transform(self, texts: list[str]) -> list[list[float]]:
        count_matrix = super().transform(texts)
        return self.transformer.transform(count_matrix)