alikia2x/cvsa
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update: use tqdm for filter model prediction

Browse Source
This commit is contained in:
alikia2x (寒寒) 2025-01-25 04:54:38 +08:00
parent 7c318c942f
commit f3db68e288
Signed by: alikia2x
GPG Key ID: 56209E0CCD8420C6
1 changed files with 3 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
filter/predict.py
View File

@ -5,6 +5,7 @@ import json
import torch import torch
from modelV3_4 import VideoClassifierV3_4 from modelV3_4 import VideoClassifierV3_4
from sentence_transformers import SentenceTransformer from sentence_transformers import SentenceTransformer
from tqdm import tqdm # 导入 tqdm
# 数据库配置 # 数据库配置
DATABASE_PATH = "./data/main.db" DATABASE_PATH = "./data/main.db"
@ -76,7 +77,7 @@ def process_entries():
# 批量处理并保存结果 # 批量处理并保存结果
with open(OUTPUT_FILE, "w", encoding="utf-8") as output: with open(OUTPUT_FILE, "w", encoding="utf-8") as output:
batch_data = [] batch_data = []
for idx, aid in enumerate(aids, 1): for idx, aid in tqdm(enumerate(aids, 1), total=len(aids), desc="处理条目"):
try: try:
# 获取并解析数据 # 获取并解析数据
raw_data = fetch_entry_data(conn, aid) raw_data = fetch_entry_data(conn, aid)
@ -111,10 +112,6 @@ def process_entries():
}, ensure_ascii=False) + "\n") }, ensure_ascii=False) + "\n")
batch_data = [] # 清空批量数据 batch_data = [] # 清空批量数据
# 进度显示
if idx % 100 == 0:
print(f"已处理 {idx}/{len(aids)} 条...")
except Exception as e: except Exception as e:
print(f"处理aid {aid} 时出错: {str(e)}") print(f"处理aid {aid} 时出错: {str(e)}")