text_data_processing.py

import re
import os
import time
import nltk
from nltk.tokenize import word_tokenize, sent_tokenize
from nltk.corpus import stopwords
from nltk.probability import FreqDist
from nltk.stem import WordNetLemmatizer
from rich.progress import Progress, TextColumn, BarColumn, TimeElapsedColumn
from data_processing_common import sanitize_filename


def summarize_text_content(text, text_inference):
    """
    对给定的文本内容进行摘要总结。

    参数:
        text: 需要总结的文本内容
        text_inference: 文本推理模型实例

    返回:
        摘要文本字符串
    """
    # 构建提示词，要求生成不超过150词的摘要
    prompt = f"""Provide a concise and accurate summary of the following text, focusing on the main ideas and key details.
Limit your summary to a maximum of 150 words.

Output only the summary, without any additional text.

Note: Please write your summary in Chinese language.
Text: {text}

Summary:
"""

    # 调用模型生成摘要
    response = text_inference.create_completion(prompt)
    # summary = response['choices'][0]['text'].strip()
    if response and 'choices' in response and response['choices']:
        summary = response['choices'][0]['text'].strip()
    else:
        summary = "Defult Summary"
        print(f"⚠️ 使用默认summary: {summary}")

    return summary


def process_single_text_file(args, text_inference, silent=False, log_file=None):
    """
    处理单个文本文件，生成元数据。

    参数:
        args: 包含文件路径和文本内容的元组
        text_inference: 文本推理模型实例
        silent: 是否静默模式（不输出到控制台）
        log_file: 日志文件路径

    返回:
        包含文件元数据的字典
    """
    file_path, text = args
    start_time = time.time()

    # 为当前文件创建进度条实例
    with Progress(
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            TimeElapsedColumn()
    ) as progress:
        # 添加处理任务到进度条
        task_id = progress.add_task(f"正在处理 {os.path.basename(file_path)}", total=1.0)
        # 生成文本元数据（文件夹名、文件名、描述）
        foldername, filename, description = generate_text_metadata(text, file_path, progress, task_id, text_inference)

    end_time = time.time()
    time_taken = end_time - start_time

    # 构建日志消息（中文输出）
    message = f"文件: {file_path}\n耗时: {time_taken:.2f} 秒\n描述: {description}\n文件夹名: {foldername}\n生成的文件名: {filename}\n"

    # 根据静默模式决定输出方式
    if silent:
        # 静默模式：写入日志文件
        if log_file:
            with open(log_file, 'a', encoding='utf-8') as f:
                f.write(message + '\n')
    else:
        # 非静默模式：输出到控制台
        print(message)

    # 返回处理结果
    return {
        'file_path': file_path,
        'foldername': foldername,
        'filename': filename,
        'description': description
    }


def process_text_files(text_tuples, text_inference, silent=False, log_file=None):
    """
    顺序处理多个文本文件。

    参数:
        text_tuples: 文本文件元组列表，每个元组包含(文件路径, 文本内容)
        text_inference: 文本推理模型实例
        silent: 是否静默模式（不输出到控制台）
        log_file: 日志文件路径

    返回:
        包含所有文件处理结果的列表
    """
    results = []
    # 遍历所有文本文件进行处理
    for args in text_tuples:
        # 处理单个文件并获取元数据
        data = process_single_text_file(args, text_inference, silent=silent, log_file=log_file)
        results.append(data)
    return results


def generate_text_metadata(input_text, file_path, progress, task_id, text_inference):
    """
    为文本文档生成描述、文件夹名称和文件名。

    参数:
        input_text: 输入的文本内容
        file_path: 文件路径
        progress: 进度条对象
        task_id: 任务ID
        text_inference: 文本推理模型

    返回:
        tuple: (文件夹名, 文件名, 描述)
    """

    # 处理文本文件的总步骤数
    total_steps = 3

    # 步骤 1: 生成文档描述
    print(f"正在生成文档描述...")
    description = summarize_text_content(input_text, text_inference)
    progress.update(task_id, advance=1 / total_steps)
    print(f"描述生成完成: {description}...")  # 只显示前100个字符

    # 步骤 2: 生成文件名
    print(f"正在生成文件名...")
    filename_prompt = f"""Based on the summary below, generate a specific and descriptive filename that captures the essence of the document.
Limit the filename to a maximum of 3 words. Use nouns and avoid starting with verbs like 'depicts', 'shows', 'presents', etc.
Do not include any data type words like 'text', 'document', 'pdf', etc. Use only letters and connect words with underscores.

Summary: {description}

Examples:
1. Summary: A research paper on the fundamentals of string theory.
   Filename: fundamentals_of_string_theory

2. Summary: An article discussing the effects of climate change on polar bears.
   Filename: climate_change_polar_bears
Note: Please write your Filename in Chinese language (中文).

Now generate the filename.

Output only the filename, without any additional text.

Filename:
"""
    filename_response = text_inference.create_completion(filename_prompt)
    # filename = filename_response['choices'][0]['text'].strip()
    if filename_response and 'choices' in filename_response and filename_response['choices']:
        filename = filename_response['choices'][0]['text'].strip()
    else:
        filename = os.path.basename(file_path)
        print(f"⚠️ 使用原文件名: {filename}")

    # 移除可能存在的 'Filename:' 前缀
    filename = re.sub(r'^Filename:\s*', '', filename, flags=re.IGNORECASE).strip()
    print(f"原始文件名: {filename}")
    progress.update(task_id, advance=1 / total_steps)

    # 步骤 3: 从摘要生成文件夹名称
    print(f"正在生成文件夹名称...")
    foldername_prompt = f"""Based on the summary below, generate a general category or theme that best represents the main subject of this document.
This will be used as the folder name. Limit the category to a maximum of 2 words. Use nouns and avoid verbs.
Do not include specific details, words from the filename, or any generic terms like 'untitled' or 'unknown'.

Summary: {description}

Examples:
1. Summary: A research paper on the fundamentals of string theory.
   Category: physics

2. Summary: An article discussing the effects of climate change on polar bears.
   Category: environment
Note: Please write your Category in Chinese language (中文).

Now generate the category.

Output only the category, without any additional text.

Category:
"""
    foldername_response = text_inference.create_completion(
        prompt=foldername_prompt,
        timeout=120
    )
    if foldername_response is None:
        print(f"[警告] 模型响应超时，使用默认分类")
        foldername = "Uncategorized"
    elif 'choices' not in foldername_response or len(foldername_response['choices']) == 0:
        print(f"[警告] 模型响应格式错误，使用默认分类")
        foldername = "Uncategorized"
    else:
        foldername = foldername_response['choices'][0]['text'].strip()

        # 同样处理 filename_response
    if filename_response is None:
        filename = os.path.basename(file_path)
    elif 'choices' not in filename_response or len(filename_response['choices']) == 0:
        filename = os.path.basename(file_path)
    else:
        filename = filename_response['choices'][0]['text'].strip()


    # foldername = foldername_response['choices'][0]['text'].strip()
    # # 移除可能存在的 'Category:' 前缀
    # foldername = re.sub(r'^Category:\s*', '', foldername, flags=re.IGNORECASE).strip()
    print(f"原始文件夹名: {foldername}")
    progress.update(task_id, advance=1 / total_steps)

    # 定义不需要的词汇和停用词
    unwanted_words_english  = set([
        'the', 'and', 'based', 'generated', 'this', 'is', 'filename', 'file', 'document', 'text', 'output', 'only',
        'below', 'category',
        'summary', 'key', 'details', 'information', 'note', 'notes', 'main', 'ideas', 'concepts', 'in', 'on', 'of',
        'with', 'by', 'for',
        'to', 'from', 'a', 'an', 'as', 'at', 'i', 'we', 'you', 'they', 'he', 'she', 'it', 'that', 'which', 'are',
        'were', 'was', 'be',
        'have', 'has', 'had', 'do', 'does', 'did', 'but', 'if', 'or', 'because', 'about', 'into', 'through', 'during',
        'before', 'after',
        'above', 'below', 'any', 'each', 'few', 'more', 'most', 'other', 'some', 'such', 'no', 'nor', 'not', 'only',
        'own', 'same', 'so',
        'than', 'too', 'very', 's', 't', 'can', 'will', 'just', 'don', 'should', 'now', 'new', 'depicts', 'show',
        'shows', 'display',
        'illustrates', 'presents', 'features', 'provides', 'covers', 'includes', 'discusses', 'demonstrates',
        'describes'
    ])
    unwanted_words_chinese  = set([
        # 原英文停用词的中文翻译
        '这个', '那个', '基于', '生成', '这是', '文件名', '文件', '文档', '文本', '输出', '仅', '只',
        '下面', '类别', '分类',
        '摘要', '总结', '关键', '详情', '细节', '信息', '笔记', '注释', '主要', '想法', '概念', '在', '于', '的',
        '与', '由', '为', '通过',
        '到', '从', '一个', '作为', '我', '我们', '你', '你们', '他', '她', '它', '那', '哪个', '是',
        '有', '做', '但', '如果', '或', '因为', '关于', '进入', '通过', '期间',
        '之前', '之后', '以前', '以后',
        '上面', '下面', '任何', '每个', '少数', '更多', '最多', '其他', '一些', '这样', '没有', '不', '仅仅',
        '自己', '相同', '所以',
        '比', '太', '非常', '可以', '将', '只是', '应该', '现在', '新', '描述', '显示',
        '展示', '说明', '呈现', '特征', '提供', '涵盖', '包括', '讨论', '演示',

        # 额外添加的常见中文停用词
        '的', '了', '在', '是', '我', '有', '和', '就', '不', '人', '都', '一', '一个', '上', '也', '很', '到', '说',
        '要', '去', '你',
        '会', '着', '没有', '看', '好', '自己', '这', '那', '里', '什么', '就是', '啊', '哦', '呢', '吧', '么', '吗',

        # 文档相关常用词
        '内容', '主题', '标题', '章节', '部分', '段落', '页面', '资料', '材料', '数据',

        # 时间相关
        '时候', '时间', '现在', '过去', '未来', '今天', '昨天', '明天', '当时', '目前',

        # 程度副词
        '非常', '很', '太', '特别', '十分', '相当', '比较', '最', '更', '极其', '颇为',

        # 方位词
        '上', '下', '左', '右', '前', '后', '里', '外', '中', '内', '东', '西', '南', '北', '旁边', '附近',

        # 连接词
        '并且', '而且', '或者', '但是', '可是', '然而', '因此', '所以', '如果', '虽然', '尽管', '无论',

        # 代词
        '他们', '她们', '它们', '咱们', '您', '自己', '本身', '彼此', '互相', '各自',

        # 量词
        '个', '只', '条', '张', '片', '块', '本', '篇', '则', '份', '项', '次', '遍', '番',

        # 助词
        '得', '地', '着', '了', '过', '起来', '下去', '出来', '进去',

        # 指示词
        '这些', '那些', '某些', '各种', '种种', '诸如', '如此', '这样', '那样',

        # 疑问词
        '什么', '怎么', '怎样', '如何', '为什么', '哪里', '哪个', '谁', '多少', '几',

        # 常见动词
        '进行', '实现', '完成', '开始', '结束', '继续', '成为', '变成', '得到', '给予',

        # 常见形容词
        '重要', '主要', '一般', '普通', '常见', '特殊', '具体', '抽象', '详细', '简单',

        # 常见介词
        '对于', '关于', '按照', '根据', '依据', '由于', '至于', '对',

        # 其他常见词
        '可能', '能够', '需要', '应该', '必须', '已经', '正在', '曾经', '总是', '经常', '有时', '偶尔',
        '大概', '也许', '或许', '似乎', '好像', '仿佛', '简直', '几乎', '差不多',
        # 文件整理相关可能出现的词
        '文件夹', '目录', '路径', '格式', '类型', '版本', '副本', '备份', '原始',

    ])

    stop_words_english = set(stopwords.words('english'))  # NLTK英文停用词
    all_unwanted_words = unwanted_words_english.union(stop_words_english).union(unwanted_words_chinese)  # 中英文全部合并

    # 词形还原器（仅用于英文）
    lemmatizer = WordNetLemmatizer()

    # 清理和处理AI输出的函数
    def clean_ai_output(text, max_words):
        """
        清理AI生成的文本输出

        参数:
            text: 待清理的文本
            max_words: 最大词数限制

        返回:
            str: 清理后用下划线连接的词汇
        """
        # 移除特殊字符和数字
        text = re.sub(r'[^\w\s]', ' ', text)
        text = re.sub(r'\d+', '', text)
        text = text.strip()
        # 分割连接的单词 (例如: 'mathOperations' -> 'math Operations')
        text = re.sub(r'([a-z])([A-Z])', r'\1 \2', text)
        # 分词和词形还原
        words = word_tokenize(text)
        words = [word.lower() for word in words if word.isalpha()]
        words = [lemmatizer.lemmatize(word) for word in words]
        # 移除不需要的词汇和重复词
        filtered_words = []
        seen = set()
        for word in words:
            if word not in all_unwanted_words and word not in seen:
                filtered_words.append(word)
                seen.add(word)
        # 限制最大词数
        filtered_words = filtered_words[:max_words]
        return '_'.join(filtered_words)

    # 处理文件名
    print(f"正在清理文件名...")
    filename = clean_ai_output(filename, max_words=3)
    if not filename or filename.lower() in ('untitled', ''):
        # 如果文件名为空或无效，从描述中提取关键词
        print(f"文件名无效，从描述中提取关键词...")
        filename = clean_ai_output(description, max_words=3)
    if not filename:
        # 如果仍然为空，使用原始文件名
        print(f"无法生成有效文件名，使用原始文件名...")
        filename = 'document_' + os.path.splitext(os.path.basename(file_path))[0]

    sanitized_filename = sanitize_filename(filename, max_words=3)
    print(f"最终文件名: {sanitized_filename}")

    # 处理文件夹名
    print(f"正在清理文件夹名称...")
    foldername = clean_ai_output(foldername, max_words=2)
    if not foldername or foldername.lower() in ('untitled', ''):
        # 如果文件夹名为空或无效，从描述中提取关键词
        print(f"文件夹名无效，从描述中提取关键词...")
        foldername = clean_ai_output(description, max_words=2)
        if not foldername:
            # 如果仍然为空，使用默认名称
            print(f"无法生成有效文件夹名，使用默认名称...")
            foldername = 'documents'

    sanitized_foldername = sanitize_filename(foldername, max_words=2)
    print(f"最终文件夹名: {sanitized_foldername}")

    print(f"元数据生成完成 - 文件夹: {sanitized_foldername}, 文件名: {sanitized_filename}")
    return sanitized_foldername, sanitized_filename, description