yifan0sun commited on
Commit
002c53a
Β·
1 Parent(s): 2be5913

Update server.py

Browse files
Files changed (1) hide show
  1. server.py +77 -0
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 ---")