Spaces:
Sleeping
Sleeping
File size: 4,350 Bytes
a8c4bc6 | 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 | import os
import html as html_lib
import re
import subprocess
import sys
from pathlib import Path
from urllib.parse import parse_qs, urlencode, urljoin, urlparse
from urllib.request import build_opener, HTTPCookieProcessor, Request
ROOT_DIR = Path(__file__).resolve().parent.parent
CHECKPOINT_DIR = Path(os.getenv("MINDFUL_CHECKPOINT_DIR", ROOT_DIR / "checkpoints"))
def extract_google_drive_id(url: str) -> str | None:
parsed = urlparse(url)
if "drive.google.com" not in parsed.netloc:
return None
query_id = parse_qs(parsed.query).get("id")
if query_id:
return query_id[0]
parts = [part for part in parsed.path.split("/") if part]
if "d" in parts:
index = parts.index("d")
if index + 1 < len(parts):
return parts[index + 1]
return None
def ensure_file_from_env(env_name: str, target_name: str) -> None:
url = os.getenv(env_name)
target_path = CHECKPOINT_DIR / target_name
if target_path.exists():
print(f"[startup] Found {target_name}, skipping download.")
return
if not url:
print(f"[startup] {target_name} missing and {env_name} is not set.", file=sys.stderr)
return
CHECKPOINT_DIR.mkdir(parents=True, exist_ok=True)
print(f"[startup] Downloading {target_name} from {env_name}...")
file_id = extract_google_drive_id(url)
if file_id:
download_google_drive_file(file_id, target_path)
else:
raise RuntimeError(f"Unsupported download URL for {env_name}. Use a Google Drive share link.")
def download_google_drive_file(file_id: str, target_path: Path) -> None:
opener = build_opener(HTTPCookieProcessor())
base_url = f"https://drive.google.com/uc?export=download&id={file_id}"
with opener.open(Request(base_url, headers={"User-Agent": "Mozilla/5.0"})) as response:
html = response.read().decode("utf-8", errors="ignore")
download_url = base_url
form_match = re.search(r'<form[^>]+id="download-form"[^>]+action="([^"]+)"', html)
if form_match:
action = html_lib.unescape(form_match.group(1))
hidden_inputs = re.findall(
r'<input[^>]+type="hidden"[^>]+name="([^"]+)"[^>]+value="([^"]*)"',
html,
)
params = {name: html_lib.unescape(value) for name, value in hidden_inputs}
params.setdefault("id", file_id)
download_url = urljoin("https://drive.google.com", action)
if params:
download_url = f"{download_url}?{urlencode(params)}"
else:
confirm_token = None
token_match = re.search(r'name="confirm"\s+value="([^"]+)"', html)
if token_match:
confirm_token = token_match.group(1)
else:
for marker in ("confirm=", "confirm=t&confirm="):
if marker in html:
fragment = html.split(marker, 1)[1]
confirm_token = fragment.split("&", 1)[0].split('"', 1)[0]
break
if confirm_token:
download_url = f"https://drive.google.com/uc?export=download&confirm={confirm_token}&id={file_id}"
with opener.open(Request(download_url, headers={"User-Agent": "Mozilla/5.0"})) as response:
content_type = response.headers.get("Content-Type", "")
if "text/html" in content_type:
body = response.read().decode("utf-8", errors="ignore")
raise RuntimeError(
"Google Drive download failed. Make sure the file is shared as "
"'Anyone with the link' and that download quota is not exceeded.\n"
f"Response snippet: {body[:300]}"
)
with open(target_path, "wb") as output_file:
while True:
chunk = response.read(1024 * 1024)
if not chunk:
break
output_file.write(chunk)
def main() -> None:
ensure_file_from_env("BEST_MODEL_URL", "best_model.pt")
ensure_file_from_env("FINAL_MODEL_URL", "final_model.pt")
port = os.getenv("PORT", "8000")
subprocess.run(
[
sys.executable,
"-m",
"uvicorn",
"backend.main:app",
"--host",
"0.0.0.0",
"--port",
port,
],
check=True,
)
if __name__ == "__main__":
main()
|