Spaces:
Sleeping
Sleeping
File size: 5,279 Bytes
f624d68 |
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 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 |
import os
import json
import sys
import time
import grequests
import sqlite3
from tqdm import tqdm
import list_files
import list_repos
SQLITE3_DB = "data/reconstructions.sqlite3"
HF_ENDPOINT = os.getenv("HF_ENDPOINT", "https://huggingface.co")
XET_CAS_ENDPOINT = os.getenv("XET_CAS_ENDPOINT", "https://cas-server.xethub.hf.co")
RESOLVE_URL_TEMPLATE = HF_ENDPOINT + "/{}/resolve/main"
def exception_handler(req, exc):
print(exc, file=sys.stderr)
def list_reconstructions_from_hub(repo):
print(
"Listing reconstructions using:\nHF Hub Endpoint: {}\nXet CAS Endpoint: {}".format(
HF_ENDPOINT, XET_CAS_ENDPOINT
),
file=sys.stderr,
)
ret = []
files = []
resolve_reqs = []
reconstruct_reqs = []
err_count = 0
print("Listing files for repo {}".format(repo), file=sys.stderr)
total = 0
for i, file in tqdm(enumerate(list_files.list_lfs_files(repo))):
total += 1
files.append(file["name"])
if repo.startswith("models/"):
repo = repo.replace("models/", "", 1)
url = file["name"].replace(repo, RESOLVE_URL_TEMPLATE.format(repo), 1)
headers = {"Authorization": "Bearer {}".format(os.getenv("HF_TOKEN"))}
resolve_reqs.append(
grequests.head(url, headers=headers, allow_redirects=False)
)
print("", file=sys.stderr)
print("Calling /resolve/ for repo {}".format(repo), file=sys.stderr)
for i, resp in tqdm(
grequests.imap_enumerated(
resolve_reqs, size=4, exception_handler=exception_handler
),
total=total,
):
if resp is None:
err_count += 1
continue
# todo: use refresh_route when access_token is expired
refresh_route = resp.headers.get("x-xet-refresh-route")
xet_hash = resp.headers.get("x-xet-hash")
access_token = resp.headers.get("x-xet-access-token")
if xet_hash is not None and xet_hash != "":
url = "{}/reconstruction/{}".format(XET_CAS_ENDPOINT, xet_hash)
headers = {"Authorization": "Bearer {}".format(access_token)}
reconstruct_reqs.append(grequests.get(url, headers=headers))
print("", file=sys.stderr)
print(
"Calling /reconstruct/ with grequests for repo {}".format(repo),
file=sys.stderr,
)
for i, resp in tqdm(
grequests.imap_enumerated(
reconstruct_reqs, size=4, exception_handler=exception_handler
),
total=total,
):
if resp is None:
continue
if resp.status_code != 200:
continue
body = resp.json()
for term in body["terms"]:
entry = {
"start": term["range"]["start"],
"end": term["range"]["end"],
"file_path": files[i + err_count],
"xorb_id": term["hash"],
"unpacked_length": term["unpacked_length"]
}
ret.append(entry)
return ret
def list_reconstructions(repos, limit=None):
ret = []
con = sqlite3.connect(SQLITE3_DB)
cur = con.cursor()
for repo in repos:
if limit is None:
res = cur.execute("SELECT * FROM reconstructions WHERE repo = '{}'".format(repo))
else:
res = cur.execute("SELECT * FROM reconstructions WHERE repo = '{}' LIMIT {}".format(repo, limit))
for row in res.fetchall():
entry = {
"xorb_id": row[1],
"last_updated_timestamp": row[2],
"repo": row[3],
"file_path": row[4],
"unpacked_length": row[5],
"start": row[6],
"end": row[7]
}
ret.append(entry)
return ret
def write_files_to_db(repo):
print("Opening database", SQLITE3_DB, file=sys.stderr)
con = sqlite3.connect(SQLITE3_DB)
cur = con.cursor()
print("Creating reconstructions table if not exists", file=sys.stderr)
cur.execute(
"CREATE TABLE IF NOT EXISTS reconstructions (id INTEGER PRIMARY KEY AUTOINCREMENT, xorb_id TEXT, last_updated_datetime INTEGER, repo TEXT, file_path TEXT, unpacked_length INTEGER, start INTEGER, end INTEGER)"
)
con.commit()
print("Deleting existing rows for repo {}".format(repo), file=sys.stderr)
cur.execute("DELETE FROM reconstructions WHERE repo = '{}'".format(repo))
con.commit()
print("Inserting rows from HFFileSystem query", file=sys.stderr)
for reconstruction in list_reconstructions_from_hub(repo):
query = "INSERT INTO reconstructions VALUES (NULL, '{}', {}, '{}', '{}', {}, {}, {})".format(
reconstruction["xorb_id"],
int(time.time()),
repo,
reconstruction["file_path"],
reconstruction["unpacked_length"],
reconstruction["start"],
reconstruction["end"]
)
cur.execute(query)
con.commit()
if __name__ == "__main__":
for repo in list_repos.list_repos():
write_files_to_db(repo)
print("Done writing to DB. Sample of 5 rows:")
json.dump(
list_reconstructions(list_repos.list_repos(), limit=5),
sys.stdout,
sort_keys=True,
indent=4,
)
|