generated from wystan_rin/template
60 lines
1.6 KiB
Python
60 lines
1.6 KiB
Python
|
#!/usr/bin/env python
|
|||
|
|
# -*- coding: UTF-8 -*-
|
|||
|
|
"""
|
|||
|
|
@Project ：recommender
|
|||
|
|
@File ：tf_idf.py
|
|||
|
|
@IDE ：PyCharm
|
|||
|
|
@Author ：rengengchen
|
|||
|
|
@Time ：2024/3/13 15:38
|
|||
|
|
"""
|
|||
|
|
import os.path
|
|||
|
|
import re
|
|||
|
|
|
|||
|
|
import jieba
|
|||
|
|
from sklearn.feature_extraction.text import CountVectorizer, TfidfTransformer
|
|||
|
|
|
|||
|
|
re_num = re.compile(r'\d+')
|
|||
|
|
|
|||
|
|
stop_words_file = os.path.join(os.path.dirname(__file__), 'ai/stop_words.txt')
|
|||
|
|
with open(stop_words_file, mode='r', encoding='utf-8') as f:
|
|||
|
|
stop_words = set(word.strip() for word in f.readlines())
|
|||
|
|
|
|||
|
|
|
|||
|
|
def filter_word(words):
|
|||
|
|
filtered_words = []
|
|||
|
|
for word in words:
|
|||
|
|
if re_num.match(word):
|
|||
|
|
continue
|
|||
|
|
if word in stop_words:
|
|||
|
|
continue
|
|||
|
|
filtered_words.append(word)
|
|||
|
|
return filtered_words
|
|||
|
|
|
|||
|
|
|
|||
|
|
def jieba_tokenize(text):
|
|||
|
|
words = jieba.lcut(text)
|
|||
|
|
words = filter_word(words)
|
|||
|
|
return words
|
|||
|
|
|
|||
|
|
|
|||
|
|
class TFIDF:
|
|||
|
|
def __init__(self, min_df, max_df, max_features):
|
|||
|
|
self.min_df = min_df
|
|||
|
|
self.max_df = max_df
|
|||
|
|
self.max_features = max_features
|
|||
|
|
|
|||
|
|
def count_words(self, corpus):
|
|||
|
|
cv = CountVectorizer(tokenizer=jieba_tokenize, min_df=self.min_df, max_df=self.max_df,
|
|||
|
|
max_features=self.max_features)
|
|||
|
|
return cv, cv.fit_transform(corpus)
|
|||
|
|
|
|||
|
|
def fit(self, corpus, get_feature_names=True):
|
|||
|
|
cv, words_matrix = self.count_words(corpus)
|
|||
|
|
|
|||
|
|
tfidf_transformer = TfidfTransformer()
|
|||
|
|
tfidf = tfidf_transformer.fit_transform(words_matrix)
|
|||
|
|
if get_feature_names:
|
|||
|
|
return cv.get_feature_names_out(), tfidf.toarray()
|
|||
|
|
# row doc column word
|
|||
|
|
return tfidf.toarray()
|