#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ 打开tophub_data.db数据库，读取表单，提取所有的类 访问本地ollama的api，修改类的名称为2-4个字，去掉中间的空格、特殊字符等字符 """ import requests import sqlite3 import re import time from loguru import logger # 配置日志 logger.add("db_modify.log", rotation="10 MB", level="INFO") class CategoryModifier: """类别修改器，用于优化数据库中的类别名称""" def __init__(self, db_path="tophub_data.db"): """ 初始化类别修改器 Args: db_path (str): 数据库路径 """ self.db_path = db_path self.ollama_url = "http://localhost:11434/api/generate" self.model = "qwen3:8b" def get_all_categories(self): """ 从数据库中获取所有唯一的类别 Returns: list: 包含所有唯一类别的列表 """ try: conn = sqlite3.connect(self.db_path) cursor = conn.cursor() cursor.execute("SELECT DISTINCT category FROM articles") categories = [row[0] for row in cursor.fetchall() if row[0]] conn.close() logger.info(f"成功获取 {len(categories)} 个唯一类别") return categories except Exception as e: logger.error(f"获取类别时出错: {e}") return [] def clean_category_name(self, category): """ 清理类别名称，移除特殊字符和多余空格 Args: category (str): 原始类别名称 Returns: str: 清理后的类别名称 """ # 移除特殊字符，只保留中文、英文和数字 cleaned = re.sub(r'[^\u4e00-\u9fa5a-zA-Z0-9]', '', category) # 移除多余的空格 cleaned = re.sub(r'\s+', '', cleaned) return cleaned def optimize_category_with_ollama(self, category): """ 使用Ollama API优化类别名称 Args: category (str): 原始类别名称 Returns: str: 优化后的类别名称 """ try: # 构造提示词 prompt = f"请将以下类别名称简化为3-6个汉字，去除空格和特殊符号，更容易理解，并保持原意：'{category}'。" + \ "例子一：'新科科技'，优化为'新质生产力'。例子二：'产设'，优化为'产品设计'。例子三：'史人'，优化为'历史人物'。" # 准备请求数据 data = { "model": self.model, "prompt": prompt, "stream": False } # 发送请求到Ollama API response = requests.post(self.ollama_url, json=data, timeout=30) response.raise_for_status() # 解析响应 result = response.json() optimized = result.get("response", "").strip() # 清理优化后的名称 optimized = self.clean_category_name(optimized) logger.info(f"类别 '{category}' 优化为 '{optimized}'") return optimized except Exception as e: logger.error(f"优化类别 '{category}' 时出错: {e}") # 如果API调用失败，返回清理后的原始名称 return self.clean_category_name(category) def update_category_in_db(self, old_category, new_category): """ 更新数据库中的类别名称 Args: old_category (str): 原始类别名称 new_category (str): 新的类别名称 """ try: conn = sqlite3.connect(self.db_path) cursor = conn.cursor() cursor.execute( "UPDATE articles SET category = ? WHERE category = ?", (new_category, old_category) ) count = cursor.rowcount conn.commit() conn.close() logger.info(f"成功更新类别 '{old_category}' 为 '{new_category}'，影响 {count} 条记录") except Exception as e: logger.error(f"更新类别 '{old_category}' 时出错: {e}") def process_all_categories(self): """ 处理所有类别 """ logger.info("开始处理所有类别...") # 获取所有类别 categories = self.get_all_categories() if not categories: logger.warning("未找到任何类别") return # 初始化进度统计 total_categories = len(categories) processed_count = 0 unchanged_count = 0 updated_count = 0 start_time = time.time() logger.info(f"总共需要处理 {total_categories} 个类别") # 处理每个类别 for i, category in enumerate(categories, 1): category_start_time = time.time() logger.info(f"处理进度: {i}/{total_categories} ({i/total_categories*100:.1f}%) - 类别: {category}") # 使用Ollama API优化类别名称 optimized_category = self.optimize_category_with_ollama(category) # 如果优化后的名称与原始名称不同，则更新数据库 if optimized_category != category: self.update_category_in_db(category, optimized_category) updated_count += 1 logger.info(f"类别 '{category}' 已更新为 '{optimized_category}'") else: unchanged_count += 1 logger.info(f"类别 '{category}' 无需更改") processed_count += 1 category_end_time = time.time() category_duration = category_end_time - category_start_time # 显示当前类别处理时间和平均处理时间 elapsed_time = time.time() - start_time avg_time_per_category = elapsed_time / processed_count estimated_remaining = avg_time_per_category * (total_categories - processed_count) logger.info(f"类别 '{category}' 处理完成，耗时: {category_duration:.2f}秒") logger.info(f"累计处理: {processed_count}/{total_categories} | " f"已更新: {updated_count} | 未更改: {unchanged_count} | " f"平均耗时: {avg_time_per_category:.2f}秒/类别 | " f"预计剩余时间: {estimated_remaining:.2f}秒") # 显示总体统计信息 total_duration = time.time() - start_time logger.info("="*60) logger.info("所有类别处理完成!") logger.info(f"总计处理类别数: {total_categories}") logger.info(f"更新类别数: {updated_count}") logger.info(f"未更改类别数: {unchanged_count}") logger.info(f"总耗时: {total_duration:.2f}秒") logger.info(f"平均每类别处理时间: {total_duration/total_categories:.2f}秒") logger.info("="*60) def main(): """主函数""" modifier = CategoryModifier() # 检查Ollama服务是否可用 try: response = requests.get("http://localhost:11434/api/tags", timeout=5) if response.status_code == 200: logger.info("Ollama服务可用") else: logger.warning("Ollama服务不可用，请确保服务已启动") return except Exception as e: logger.warning(f"无法连接到Ollama服务: {e}") logger.info("请确保Ollama服务已在本地运行") return # 处理所有类别 modifier.process_all_categories() if __name__ == "__main__": main()