165 lines
7.9 KiB
Python
165 lines
7.9 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
种子数据执行入口：初始化 SQLite 表结构 → 灌入模拟数据 → 构建 Chroma 向量索引。
|
||
幂等设计：重复执行会自动清空重灌，方便数据更新后一键重建。
|
||
用法：python3 seed.py [--rebuild-vector] （--rebuild-vector 强制重建向量集合）
|
||
"""
|
||
import argparse
|
||
import logging
|
||
import sys
|
||
import time
|
||
|
||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||
log = logging.getLogger("seed")
|
||
|
||
from db import init_db, execute, executemany, query_one, table_count # noqa: E402
|
||
from seed_teams import SPORTS, LEAGUES, TEAMS, STANDINGS # noqa: E402
|
||
from seed_players_a import PLAYERS_A # noqa: E402
|
||
from seed_players_b import PLAYERS_B # noqa: E402
|
||
from seed_games import GAMES, BOX # noqa: E402
|
||
from seed_news import NEWS, WIKI # noqa: E402
|
||
from seed_persons import PERSONS # noqa: E402
|
||
import vector_store # noqa: E402
|
||
|
||
|
||
def wipe():
|
||
for t in ("game_player_stats", "games", "standings", "players", "teams",
|
||
"leagues", "sports", "news", "persons"):
|
||
execute(f"DELETE FROM {t}")
|
||
execute(f"DELETE FROM sqlite_sequence WHERE name='{t}'")
|
||
log.info("已清空旧数据")
|
||
|
||
|
||
def seed_core():
|
||
"""sports / leagues / teams / players / standings / games / persons"""
|
||
for code, name, name_en in SPORTS:
|
||
execute("INSERT INTO sports(code,name,name_en) VALUES(?,?,?)", (code, name, name_en))
|
||
sport_id = query_one("SELECT id FROM sports WHERE code='basketball'")["id"]
|
||
for code, name, name_en, country, season in LEAGUES:
|
||
execute("INSERT INTO leagues(sport_id,code,name,name_en,country,season) VALUES(?,?,?,?,?,?)",
|
||
(sport_id, code, name, name_en, country, season))
|
||
league_id = query_one("SELECT id FROM leagues WHERE code='NBA'")["id"]
|
||
|
||
# 球队
|
||
for code, name, en, city, arena, founded, champs, coach, intro in TEAMS:
|
||
execute("""INSERT INTO teams(league_id,code,name,name_en,city,arena,founded,champion_count,head_coach,intro)
|
||
VALUES(?,?,?,?,?,?,?,?,?,?)""",
|
||
(league_id, code, name, en, city, arena, founded, champs, coach, intro))
|
||
team_id = {r["code"]: r["id"] for r in __import__("db").query("SELECT id,code FROM teams")}
|
||
|
||
# 球员
|
||
p_rows = []
|
||
for p in PLAYERS_A + PLAYERS_B:
|
||
(t, name, en, pos, num, h, w, country, dy, dp, sal,
|
||
pts, reb, ast, stl, blk, mn, cpts, creb, cast, cg, awards, bio) = p
|
||
p_rows.append((team_id[t], name, en, pos, num, h, w, country, dy, dp, sal,
|
||
pts, reb, ast, stl, blk, mn, cpts, creb, cast, cg, awards, bio))
|
||
executemany("""INSERT INTO players(team_id,name,name_en,position,number,height_cm,weight_kg,country,
|
||
draft_year,draft_pick,salary_m,season_pts,season_reb,season_ast,season_stl,season_blk,season_min,
|
||
career_pts,career_reb,career_ast,career_games,awards,bio)
|
||
VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)""", p_rows)
|
||
log.info("球员 %d 名", len(p_rows))
|
||
|
||
# 排名
|
||
s_rows = [(league_id, "2025-26", team_id[c], w, l, round(w / (w + l), 4), conf, rk)
|
||
for c, conf, rk, w, l in STANDINGS]
|
||
executemany("""INSERT INTO standings(league_id,season,team_id,wins,losses,win_pct,conference,rank)
|
||
VALUES(?,?,?,?,?,?,?,?)""", s_rows)
|
||
|
||
# 比赛
|
||
game_id = {}
|
||
g_rows = []
|
||
for rnd, gt, status, home, away, hs, as_, venue, bcast in GAMES:
|
||
g_rows.append((league_id, "2025-26", rnd, gt, status, team_id[home], team_id[away],
|
||
hs, as_, venue, bcast))
|
||
executemany("""INSERT INTO games(league_id,season,round_name,game_time,status,home_team_id,away_team_id,
|
||
home_score,away_score,venue,broadcast) VALUES(?,?,?,?,?,?,?,?,?,?,?)""", g_rows)
|
||
for i, (rnd, gt, *_rest) in enumerate(GAMES):
|
||
game_id[(rnd, gt)] = query_one(
|
||
"SELECT id FROM games WHERE round_name=? AND game_time=?", (rnd, gt))["id"]
|
||
|
||
# 球员技术统计（box score）
|
||
box_rows = []
|
||
for (rnd, gt), lines in BOX.items():
|
||
gid = game_id.get((rnd, gt))
|
||
if not gid:
|
||
continue
|
||
for pname, pts, reb, ast, stl, blk, mn in lines:
|
||
player = query_one("SELECT id,team_id FROM players WHERE name=?", (pname,))
|
||
if player:
|
||
box_rows.append((gid, player["id"], player["team_id"], pts, reb, ast, stl, blk, mn))
|
||
executemany("""INSERT INTO game_player_stats(game_id,player_id,team_id,points,rebounds,assists,steals,blocks,minutes)
|
||
VALUES(?,?,?,?,?,?,?,?,?)""", box_rows)
|
||
log.info("比赛 %d 场，技术统计 %d 条", len(GAMES), len(box_rows))
|
||
|
||
# 人物
|
||
per_rows = []
|
||
for t, name, en, role, role_cn, title, bio, ach in PERSONS:
|
||
per_rows.append((sport_id, league_id, team_id.get(t), name, en, role, role_cn, title, bio, ach))
|
||
executemany("""INSERT INTO persons(sport_id,league_id,team_id,name,name_en,role,role_cn,title,bio,achievements)
|
||
VALUES(?,?,?,?,?,?,?,?,?,?)""", per_rows)
|
||
log.info("人物 %d 名", len(per_rows))
|
||
|
||
# 新闻 / 百科
|
||
n_rows = []
|
||
for t, title, content, author, source, pt, tags, kind in NEWS + WIKI:
|
||
n_rows.append((sport_id, league_id, team_id.get(t), title, content, author, source, pt, tags, kind))
|
||
executemany("""INSERT INTO news(sport_id,league_id,team_id,title,content,author,source,publish_time,tags,kind)
|
||
VALUES(?,?,?,?,?,?,?,?,?,?)""", n_rows)
|
||
log.info("新闻/百科 %d 篇", len(n_rows))
|
||
|
||
|
||
def seed_vector(rebuild=False):
|
||
"""把新闻/百科分块后写入 Chroma（title + 正文 按 400 字切块，重叠 60 字）"""
|
||
from db import query as dbq
|
||
rows = dbq("SELECT id,title,content,kind,team_id,publish_time,tags FROM news ORDER BY id")
|
||
ids, docs, metas = [], [], []
|
||
for r in rows:
|
||
text = f"{r['title']}\n{r['content']}"
|
||
chunk_size, overlap = 400, 60
|
||
start = 0
|
||
while start < len(text):
|
||
chunk = text[start:start + chunk_size]
|
||
ids.append(f"news_{r['id']}_{start}")
|
||
docs.append(chunk)
|
||
metas.append({"news_id": r["id"], "title": r["title"], "kind": r["kind"],
|
||
"team_id": r["team_id"] or 0, "publish_time": r["publish_time"],
|
||
"tags": r["tags"] or ""})
|
||
start += chunk_size - overlap
|
||
log.info("向量文档 %d 块（%d 篇）", len(ids), len(rows))
|
||
if rebuild:
|
||
vector_store.reset_collection()
|
||
# 若已有数据则跳过（增量）
|
||
if vector_store.collection_count() >= len(ids):
|
||
log.info("向量库已有 %d 条，跳过写入", vector_store.collection_count())
|
||
return
|
||
if vector_store.collection_count() > 0:
|
||
vector_store.reset_collection()
|
||
for i in range(0, len(ids), 50): # 分批写入，避免单次请求过大
|
||
vector_store.add_documents(ids[i:i + 50], docs[i:i + 50], metas[i:i + 50])
|
||
log.info("向量索引完成，共 %d 条", vector_store.collection_count())
|
||
|
||
|
||
def main():
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument("--rebuild-vector", action="store_true", help="强制重建向量集合")
|
||
ap.add_argument("--no-vector", action="store_true", help="跳过向量索引")
|
||
args = ap.parse_args()
|
||
|
||
t0 = time.time()
|
||
init_db()
|
||
wipe()
|
||
seed_core()
|
||
if not args.no_vector:
|
||
seed_vector(rebuild=args.rebuild_vector)
|
||
else:
|
||
log.info("跳过向量索引")
|
||
log.info("数据灌入完成，耗时 %.1fs", time.time() - t0)
|
||
for t in ("sports", "leagues", "teams", "players", "games", "standings",
|
||
"game_player_stats", "news", "persons"):
|
||
log.info(" %-16s %d", t, table_count(t))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|