#文本特征工程详解：TF-IDF算法、相似度计算与词袋模型演进及PyTorch实现

from sklearn.feature_extraction.text import CountVectorizer
import jieba

# 定义中文分词器
def zh_tokenizer(text):
    return jieba.lcut(text)

# 示例语料
docs = [
    "自然语言处理是人工智能的重要分支",
    "机器学习和深度学习是核心技术",
    "自然语言处理和机器学习密切相关"
]

# 初始化词袋向量化器
bow_vec = CountVectorizer(
    tokenizer=zh_tokenizer,
    lowercase=False,  # 中文不需要
    token_pattern=None  # 用自定义分词器
)

# 拟合语料，生成矩阵
bow_matrix = bow_vec.fit_transform(docs)

print("词汇表:", bow_vec.get_feature_names_out())
print("词袋矩阵:\n", bow_matrix.toarray())

from sklearn.feature_extraction.text import TfidfVectorizer
import jieba
import numpy as np

def zh_tokenizer(text):
    return jieba.lcut(text)

# 更丰富的语料
docs = [
    "自然语言处理是人工智能的重要分支",
    "机器学习和深度学习是人工智能的核心技术",
    "深度学习在计算机视觉领域应用广泛",
    "数据科学结合统计学和计算机科学"
]

# 高级向量化器配置
tfidf_vec = TfidfVectorizer(
    tokenizer=zh_tokenizer,
    max_features=1000,  # 限制词汇表大小，降维
    min_df=1,           # 至少在1个文档出现
    max_df=0.8,         # 过滤80%以上文档都有的词（停用词）
    ngram_range=(1,2),  # 同时考虑单字/词（1-gram）和双词（2-gram，比如“自然语言”）
    norm='l2',          # L2归一化，便于计算余弦相似度
    sublinear_tf=True   # 用对数缩放TF，避免词频过高的词权重过大
)

# 生成矩阵
tfidf_matrix = tfidf_vec.fit_transform(docs)

# 查看文档1的Top3关键词
feature_names = tfidf_vec.get_feature_names_out()
doc1_tfidf = tfidf_matrix[0].toarray().flatten()
top3_idx = np.argsort(doc1_tfidf)[::-1][:3]
print("文档1的Top3关键词：")
for idx in top3_idx:
    print(f"  {feature_names[idx]}: {doc1_tfidf[idx]:.4f}")

from sklearn.metrics.pairwise import cosine_similarity

# 计算整个语料的相似度矩阵
sim_matrix = cosine_similarity(tfidf_matrix)
print("语料相似度矩阵：\n", sim_matrix.round(4))

import jieba
import torch
from collections import defaultdict

def zh_tokenizer(text):
    return jieba.lcut(text)

class PyTorchTFIDF:
    def __init__(self, max_features=1000, min_df=1, max_df=0.8):
        self.max_features = max_features
        self.min_df = min_df
        self.max_df = max_df
        self.vocab = None
        self.idf = None
    
    def fit(self, texts):
        # 分词
        tokenized = [zh_tokenizer(t) for t in texts]
        total_docs = len(tokenized)
        
        # 统计词频和文档频率
        word_doc_count = defaultdict(int)
        word_total_count = defaultdict(int)
        for tokens in tokenized:
            unique_tokens = set(tokens)
            for token in unique_tokens:
                word_doc_count[token] += 1
            for token in tokens:
                word_total_count[token] += 1
        
        # 过滤词汇
        filtered_words = [
            w for w in word_total_count 
            if self.min_df <= word_doc_count[w] <= self.max_df * total_docs
        ]
        # 取TopN词频
        filtered_words.sort(key=lambda x: word_total_count[x], reverse=True)
        self.vocab = {w: i for i, w in enumerate(filtered_words[:self.max_features])}
        
        # 计算IDF
        self.idf = torch.zeros(len(self.vocab))
        for w, idx in self.vocab.items():
            # 平滑处理（避免除以0）
            self.idf[idx] = torch.log(torch.tensor(total_docs / (word_doc_count[w] + 1))) + 1
    
    def transform(self, texts):
        tokenized = [zh_tokenizer(t) for t in texts]
        tfidf_matrix = torch.zeros(len(texts), len(self.vocab))
        
        # 计算TF
        for i, tokens in enumerate(tokenized):
            token_count = defaultdict(int)
            for t in tokens:
                if t in self.vocab:
                    token_count[t] += 1
            # 归一化TF
            doc_len = len(tokens) if len(tokens) > 0 else 1
            for t, cnt in token_count.items():
                tf = cnt / doc_len
                tfidf_matrix[i, self.vocab[t]] = tf * self.idf[self.vocab[t]]
        
        # L2归一化
        norm = torch.norm(tfidf_matrix, p=2, dim=1, keepdim=True)
        norm[norm == 0] = 1  # 避免空文档除以0
        tfidf_matrix /= norm
        return tfidf_matrix

# 测试
pytorch_tfidf = PyTorchTFIDF()
pytorch_tfidf.fit(docs)
pytorch_matrix = pytorch_tfidf.transform(docs)
print("PyTorch TF-IDF矩阵形状:", pytorch_matrix.shape)

class SimpleDocSearch:
    def __init__(self, docs):
        self.docs = docs
        # 用Sklearn的更稳定
        self.vec = TfidfVectorizer(tokenizer=zh_tokenizer, ngram_range=(1,2), norm='l2')
        self.matrix = self.vec.fit_transform(docs)
    
    def search(self, query, top_k=2):
        query_vec = self.vec.transform([query])
        sims = cosine_similarity(query_vec, self.matrix).flatten()
        top_idx = sims.argsort()[::-1][:top_k]
        return [(self.docs[i], sims[i].round(4)) for i in top_idx]

# 测试
searcher = SimpleDocSearch(docs)
print("搜索结果（查询：计算机）：", searcher.search("计算机"))