168 lines
7.1 KiB
Python
168 lines
7.1 KiB
Python
|
import os
|
|||
|
|
import json
|
|||
|
|
import pymysql
|
|||
|
|
from pymysql.cursors import DictCursor
|
|||
|
|
import pymongo
|
|||
|
|
import requests
|
|||
|
|
from dotenv import load_dotenv
|
|||
|
|
from typing import List, Dict, Generator
|
|||
|
|
|
|||
|
|
# 加载环境变量
|
|||
|
|
load_dotenv()
|
|||
|
|
|
|||
|
|
class StudentDataProcessorLMStudio:
|
|||
|
|
def __init__(self):
|
|||
|
|
# 初始化 MySQL 连接（支持批量遍历）
|
|||
|
|
self.mysql_conn = pymysql.connect(
|
|||
|
|
host=os.getenv("MYSQL_HOST"),
|
|||
|
|
port=int(os.getenv("MYSQL_PORT")),
|
|||
|
|
user=os.getenv("MYSQL_USER"),
|
|||
|
|
password=os.getenv("MYSQL_PASSWORD"),
|
|||
|
|
database=os.getenv("MYSQL_DB"),
|
|||
|
|
charset="utf8mb4"
|
|||
|
|
)
|
|||
|
|
self.mysql_cursor = self.mysql_conn.cursor()
|
|||
|
|
|
|||
|
|
# 初始化 MongoDB 连接
|
|||
|
|
self.mongo_client = pymongo.MongoClient(os.getenv("MONGO_URI"))
|
|||
|
|
self.mongo_db = self.mongo_client[os.getenv("MONGO_DB")]
|
|||
|
|
self.mongo_collection = self.mongo_db[os.getenv("MONGO_COLLECTION")]
|
|||
|
|
|
|||
|
|
# LM Studio 配置
|
|||
|
|
self.lm_studio_api_url = os.getenv("LM_STUDIO_API_URL")
|
|||
|
|
self.lm_studio_model = os.getenv("LM_STUDIO_MODEL_NAME")
|
|||
|
|
self.lm_studio_api_token = os.getenv("LM_STUDIO_API_TOKEN")
|
|||
|
|
self.batch_size = int(os.getenv("BATCH_SIZE"))
|
|||
|
|
|
|||
|
|
def batch_read_student_data(self) -> Generator[List[Dict], None, None]:
|
|||
|
|
"""分批遍历 MySQL 所有学生数据（避免一次性加载过多数据）"""
|
|||
|
|
try:
|
|||
|
|
# 先获取总条数，用于分批
|
|||
|
|
count_query = "SELECT COUNT(*) as total FROM student WHERE is_deleted = 0"
|
|||
|
|
self.mysql_cursor.execute(count_query)
|
|||
|
|
total = self.mysql_cursor.fetchone()["total"]
|
|||
|
|
print(f"MySQL 中共计 {total} 条学生数据，将按每批 {self.batch_size} 条处理")
|
|||
|
|
|
|||
|
|
# 分批读取
|
|||
|
|
offset = 0
|
|||
|
|
while offset < total:
|
|||
|
|
query = f"""
|
|||
|
|
SELECT id, name, age, gender, class, phone, email, address
|
|||
|
|
FROM student
|
|||
|
|
WHERE is_deleted = 0
|
|||
|
|
LIMIT {self.batch_size} OFFSET {offset}
|
|||
|
|
"""
|
|||
|
|
self.mysql_cursor.execute(query)
|
|||
|
|
batch_data = self.mysql_cursor.fetchall()
|
|||
|
|
if not batch_data:
|
|||
|
|
break
|
|||
|
|
yield batch_data
|
|||
|
|
offset += self.batch_size
|
|||
|
|
print(f"已读取第 {offset//self.batch_size} 批数据（累计 {offset} 条）")
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"分批读取 MySQL 数据失败：{str(e)}")
|
|||
|
|
raise
|
|||
|
|
|
|||
|
|
def process_batch_with_lmstudio(self, batch_data: List[Dict]) -> List[Dict]:
|
|||
|
|
"""调用 LM Studio 本地模型整理单批数据"""
|
|||
|
|
if not batch_data:
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
# 构造适配本地模型的提示词（更简洁，避免本地模型处理复杂指令出错）
|
|||
|
|
prompt = f"""
|
|||
|
|
请严格按照以下要求整理学生数据，仅返回JSON数组，不要添加任何额外文字、解释或备注：
|
|||
|
|
整理规则：
|
|||
|
|
1. 标准化字段：
|
|||
|
|
- 姓名：去除首尾空格，仅保留中文（如" 李四 "→"李四"）
|
|||
|
|
- 年龄：转为整数，非数字则设为null
|
|||
|
|
- 手机号：仅保留11位纯数字，不符合则设为null
|
|||
|
|
- 邮箱：验证格式，无效则设为null
|
|||
|
|
- 性别：非"男"/"女"则设为"未知"
|
|||
|
|
- 班级：为空则设为"未分配"
|
|||
|
|
2. 新增字段：
|
|||
|
|
- student_id：格式为"2024"+补零后的id（如id=5→"202400005"，id=123→"202400123"）
|
|||
|
|
- data_quality："完整"（所有字段非null）/"部分缺失"（1-2个字段null）/"严重缺失"（≥3个字段null）
|
|||
|
|
原始数据：{json.dumps(batch_data, ensure_ascii=False, indent=2)}
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
# 构造 LM Studio API 请求体（兼容 OpenAI 格式）
|
|||
|
|
payload = {
|
|||
|
|
"model": self.lm_studio_model,
|
|||
|
|
"messages": [
|
|||
|
|
{"role": "system", "content": "你是数据整理专家，仅返回JSON格式的处理结果，无其他内容"},
|
|||
|
|
{"role": "user", "content": prompt}
|
|||
|
|
],
|
|||
|
|
"temperature": 0.1, # 低随机性保证结果稳定
|
|||
|
|
"max_tokens": 2000 # 足够容纳批量数据的JSON
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
# 调用 LM Studio 本地 API
|
|||
|
|
response = requests.post(
|
|||
|
|
self.lm_studio_api_url,
|
|||
|
|
headers={"Authorization": f"Bearer {self.lm_studio_api_token}"},
|
|||
|
|
json=payload,
|
|||
|
|
timeout=60 # 本地模型处理稍慢，延长超时时间
|
|||
|
|
)
|
|||
|
|
response.raise_for_status() # 抛出HTTP错误
|
|||
|
|
result = response.json()
|
|||
|
|
|
|||
|
|
# 解析本地模型返回的结果
|
|||
|
|
processed_text = result["choices"][0]["message"]["content"].strip()
|
|||
|
|
# 清理可能的多余字符（本地模型可能返回```json开头/结尾）
|
|||
|
|
if processed_text.startswith("```json"):
|
|||
|
|
processed_text = processed_text.replace("```json", "").replace("```", "").strip()
|
|||
|
|
processed_data = json.loads(processed_text)
|
|||
|
|
return processed_data
|
|||
|
|
except json.JSONDecodeError:
|
|||
|
|
print(f"本地模型返回非标准JSON，跳过该批：{processed_text[:200]}")
|
|||
|
|
return []
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"调用LM Studio失败：{str(e)}")
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
def save_to_mongodb(self, processed_data: List[Dict]):
|
|||
|
|
"""将整理后的数据存入MongoDB（支持增量插入）"""
|
|||
|
|
if not processed_data:
|
|||
|
|
return
|
|||
|
|
try:
|
|||
|
|
# 批量插入，跳过重复id（避免重复存储）
|
|||
|
|
for data in processed_data:
|
|||
|
|
self.mongo_collection.update_one(
|
|||
|
|
{"id": data["id"]}, # 以原始id为唯一键
|
|||
|
|
{"$set": data}, # 存在则更新，不存在则插入
|
|||
|
|
upsert=True
|
|||
|
|
)
|
|||
|
|
print(f"成功存入/更新 {len(processed_data)} 条数据到MongoDB")
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"存入MongoDB失败：{str(e)}")
|
|||
|
|
|
|||
|
|
def close_connections(self):
|
|||
|
|
"""关闭数据库连接"""
|
|||
|
|
if self.mysql_cursor:
|
|||
|
|
self.mysql_cursor.close()
|
|||
|
|
if self.mysql_conn:
|
|||
|
|
self.mysql_conn.close()
|
|||
|
|
if self.mongo_client:
|
|||
|
|
self.mongo_client.close()
|
|||
|
|
print("所有连接已关闭")
|
|||
|
|
|
|||
|
|
def run(self):
|
|||
|
|
"""执行完整流程：分批读取→本地模型处理→存入MongoDB"""
|
|||
|
|
try:
|
|||
|
|
# 遍历所有批次数据
|
|||
|
|
for batch in self.batch_read_student_data():
|
|||
|
|
# 处理单批数据
|
|||
|
|
processed_batch = self.process_batch_with_lmstudio(batch)
|
|||
|
|
# 存入MongoDB
|
|||
|
|
self.save_to_mongodb(processed_batch)
|
|||
|
|
print("✅ 全量数据处理完成！")
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"程序执行异常：{str(e)}")
|
|||
|
|
finally:
|
|||
|
|
self.close_connections()
|
|||
|
|
|
|||
|
|
# 程序入口
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
processor = StudentDataProcessorLMStudio()
|
|||
|
|
processor.run()
|