Pandas文本数据关键词概率分类与标签生成教程

import re
import pandas as pd
from collections import Counter

# 定义关键词类别及其对应的关键词列表
labels = {
    'fruits': ['mango', 'apple', 'lichi'],
    'animals': ['dog', 'cat', 'cow', 'monkey'],
    'country': ['us', 'ca', 'au', 'br'],
}

# 示例DataFrame
data = {
    'content': [
        'My favorite fruit is mango. I like lichies too. I live in au. Cows are domestic animals.',
        'I own RTX 4090...',
        'There is political conflict between us and ca.',
        'au, br mango, lichi apple,.... \n cat, cow, monkey donkey dogs'
    ]
}
df = pd.DataFrame(data)

print("原始DataFrame:")
print(df)

def calculate_probability(text, labels_dict):
    """
    计算给定文本中各关键词类别的概率，并返回最高概率的类别标签。

    参数:
    text (str): 输入的文本字符串。
    labels_dict (dict): 包含关键词类别的字典，键为类别名称，值为关键词列表。

    返回:
    str: 概率最高的类别标签，如果所有类别概率为0则返回 'NaN'。
    """
    # 1. 文本预处理与分词：转换为小写并提取所有单词
    words = re.findall(r'\b\w+\b', text.lower())
    word_count = len(words)

    # 如果文本为空或没有单词，则无法计算概率
    if word_count == 0:
        return 'NaN'

    # 2. 词频统计：使用Counter高效统计单词频率
    counts = Counter(words)

    # 3. 类别概率计算：遍历每个类别，累加关键词词频
    probs = {}
    for k, l_list in labels_dict.items():
        # 对于每个类别，统计其关键词在文本中出现的总次数
        # 注意：这里是精确匹配分词后的单词。
        category_keyword_count = sum(counts[w] for w in l_list)
        probs[k] = category_keyword_count / word_count

    # 4. 最高概率类别选择
    # 找出概率最高的类别名称
    max_label = max(probs, key=probs.get)

    # 如果最高概率大于0，则返回该标签；否则返回'NaN'
    return max_label if probs[max_label] > 0 else 'NaN'

# 将函数应用到DataFrame的 'content' 列
# 注意：当对Series应用函数时，函数应接收Series的单个元素作为输入。
# 额外的参数 (labels_dict) 通过关键字参数传递给apply。
df['label'] = df['content'].apply(calculate_probability, labels_dict=labels)

print("\n处理后的DataFrame:")
print(df)