File size: 4,009 Bytes
248c4e8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | import os
import sqlite3
import argparse
from pathlib import Path
from huggingface_hub import HfApi, login
def get_db_path():
# Priority list of where the DB might be
paths = [
"/data/tracker.db",
"/home/adhyansh/datasets/yt/hf_space_host/db/tracker.db",
"/home/adhyansh/datasets/yt/scripts/db/tracker.db"
]
for p in paths:
if os.path.exists(p):
return p
return None
def load_hf_token():
token = os.environ.get("HF_TOKEN", "")
if not token:
p = Path.home() / ".cache" / "huggingface" / "token"
if p.exists():
token = p.read_text().strip()
return token
def main():
parser = argparse.ArgumentParser(description="List files in an HF Repo and mark matching videos as 'done' in the local database.")
parser.add_argument("--repo-id", default="AdhyanshVerma/YT", help="Hugging Face Repository ID")
parser.add_argument("--db-path", default=None, help="Explicit path to the tracker.db SQLite database")
args = parser.parse_args()
# 1. Connect to Database
db_path = args.db_path or get_db_path()
if not db_path:
print("❌ Error: Could not find tracker.db in default locations. Please specify with --db-path")
return
print(f"✅ Using Database: {db_path}")
conn = sqlite3.connect(db_path)
cursor = conn.cursor()
# Fetch all video_ids that are NOT done
try:
cursor.execute("SELECT video_id FROM videos WHERE status != 'done'")
pending_videos = [row[0] for row in cursor.fetchall()]
print(f"📊 Found {len(pending_videos)} videos in DB that are not marked as 'done'.")
except Exception as e:
print(f"❌ Database error: {e}")
return
if not pending_videos:
print("✅ All videos are already marked as done. Exiting.")
return
# 2. Authenticate with Hugging Face
token = load_hf_token()
if token:
login(token=token)
print("✅ Authenticated with Hugging Face token.")
else:
print("⚠️ Warning: No HF_TOKEN found. If the repo is private, this might fail.")
api = HfApi()
# 3. List all files from the target repository
repo_types_to_try = ["dataset", "space", "model"]
files = []
print(f"🔍 Searching for files in repository '{args.repo_id}'...")
for repo_type in repo_types_to_try:
try:
# We use recursive=True to list ALL files without downloading them
files = api.list_repo_files(repo_id=args.repo_id, repo_type=repo_type)
print(f"✅ Successfully listed {len(files)} files as a '{repo_type}' repo.")
break
except Exception as e:
# Continue checking other repo types
pass
if not files:
print(f"❌ Error: Could not find or list files in repository '{args.repo_id}'. Make sure it exists and you have access.")
return
# 4. Check files against database and modify DB
print("\n🔄 Analyzing files against the database...")
marked_done_count = 0
for vid in pending_videos:
match_found = False
# Search for the video ID anywhere in the full file paths (this catches directory names like staging/.../Jb9Jx0KkJl8/...)
for f in files:
if vid in f:
match_found = True
break
if match_found:
print(f" 🟢 Match found for video: {vid} - Marking as 'done'!")
cursor.execute("UPDATE videos SET status='done' WHERE video_id=?", (vid,))
marked_done_count += 1
# 5. Save changes
if marked_done_count > 0:
conn.commit()
print(f"\n🎉 Success! Updated {marked_done_count} videos to 'done' in the database.")
else:
print("\n⚠️ No new matching videos found in the repository. No database changes made.")
conn.close()
if __name__ == "__main__":
main()
|