164 lines
7.9 KiB
Python
164 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())
|