Spaces:
Sleeping
Sleeping
Update server.py
Browse files
server.py
CHANGED
@@ -9,6 +9,11 @@ from ROBERTAmodel import *
|
|
9 |
from BERTmodel import *
|
10 |
from DISTILLBERTmodel import *
|
11 |
|
|
|
|
|
|
|
|
|
|
|
12 |
VISUALIZER_CLASSES = {
|
13 |
"BERT": BERTVisualizer,
|
14 |
"RoBERTa": RoBERTaVisualizer,
|
@@ -57,6 +62,78 @@ def ping():
|
|
57 |
return {"message": "pong"}
|
58 |
|
59 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
60 |
@app.post("/load_model")
|
61 |
def load_model(req: LoadModelRequest):
|
62 |
print(f"\n--- /load_model request received ---")
|
|
|
9 |
from BERTmodel import *
|
10 |
from DISTILLBERTmodel import *
|
11 |
|
12 |
+
import os
|
13 |
+
import zipfile
|
14 |
+
import shutil
|
15 |
+
|
16 |
+
|
17 |
VISUALIZER_CLASSES = {
|
18 |
"BERT": BERTVisualizer,
|
19 |
"RoBERTa": RoBERTaVisualizer,
|
|
|
62 |
return {"message": "pong"}
|
63 |
|
64 |
|
65 |
+
|
66 |
+
|
67 |
+
def zip_if_needed(src_dir, zip_path):
|
68 |
+
if os.path.exists(zip_path):
|
69 |
+
return # already zipped
|
70 |
+
print(f"π¦ Zipping {src_dir} β {zip_path}")
|
71 |
+
with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zf:
|
72 |
+
for root, _, files in os.walk(src_dir):
|
73 |
+
for file in files:
|
74 |
+
full_path = os.path.join(root, file)
|
75 |
+
rel_path = os.path.relpath(full_path, src_dir)
|
76 |
+
zf.write(full_path, arcname=rel_path)
|
77 |
+
print(f"β
Created zip: {zip_path}")
|
78 |
+
|
79 |
+
def extract_zip_if_needed(zip_path, dest_dir):
|
80 |
+
if os.path.exists(dest_dir):
|
81 |
+
print(f"β
Already exists: {dest_dir}")
|
82 |
+
return
|
83 |
+
print(f"π Extracting {zip_path} β {dest_dir}")
|
84 |
+
with zipfile.ZipFile(zip_path, 'r') as zip_ref:
|
85 |
+
zip_ref.extractall(dest_dir)
|
86 |
+
print(f"β
Extracted to: {dest_dir}")
|
87 |
+
|
88 |
+
def print_directory_tree(path):
|
89 |
+
for root, dirs, files in os.walk(path):
|
90 |
+
indent = ' ' * (root.count(os.sep) - path.count(os.sep))
|
91 |
+
print(f"{indent}π {os.path.basename(root)}/")
|
92 |
+
for f in files:
|
93 |
+
print(f"{indent} π {f}")
|
94 |
+
@app.get("/copy_and_extract")
|
95 |
+
def copy_zip_extract_and_report():
|
96 |
+
src_base = "./hf_cache"
|
97 |
+
dst_base = "/data/hf_cache"
|
98 |
+
|
99 |
+
for category in ["models", "tokenizers"]:
|
100 |
+
src_dir = os.path.join(src_base, category)
|
101 |
+
dst_dir = os.path.join(dst_base, category)
|
102 |
+
|
103 |
+
if not os.path.exists(src_dir):
|
104 |
+
continue
|
105 |
+
|
106 |
+
os.makedirs(dst_dir, exist_ok=True)
|
107 |
+
|
108 |
+
for name in os.listdir(src_dir):
|
109 |
+
full_path = os.path.join(src_dir, name)
|
110 |
+
if not os.path.isdir(full_path):
|
111 |
+
continue
|
112 |
+
|
113 |
+
zip_name = f"{name}.zip"
|
114 |
+
local_zip = os.path.join(src_dir, zip_name)
|
115 |
+
dst_zip = os.path.join(dst_dir, zip_name)
|
116 |
+
extract_path = os.path.join(dst_dir, name)
|
117 |
+
|
118 |
+
zip_if_needed(full_path, local_zip)
|
119 |
+
|
120 |
+
# Copy zip to /data
|
121 |
+
if not os.path.exists(dst_zip):
|
122 |
+
shutil.copy(local_zip, dst_zip)
|
123 |
+
print(f"π€ Copied zip to: {dst_zip}")
|
124 |
+
|
125 |
+
extract_zip_if_needed(dst_zip, extract_path)
|
126 |
+
|
127 |
+
print("\nπ¦ Local hf_cache structure:")
|
128 |
+
print_directory_tree("./hf_cache")
|
129 |
+
|
130 |
+
print("\nπΎ Persistent /data/hf_cache structure:")
|
131 |
+
print_directory_tree("/data/hf_cache")
|
132 |
+
|
133 |
+
|
134 |
+
|
135 |
+
|
136 |
+
|
137 |
@app.post("/load_model")
|
138 |
def load_model(req: LoadModelRequest):
|
139 |
print(f"\n--- /load_model request received ---")
|