File size: 3,864 Bytes
7a0c684 |
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 |
"""
Download and store HuggingFace models in device DB
"""
from typing import Optional, Dict, Any
import os
import json
from huggingface_hub import snapshot_download
import duckdb
import numpy as np
# Initialize HuggingFace token from environment
HF_TOKEN = os.getenv("HF_TOKEN")
def download_model(model_id: str, cache_dir: Optional[str] = None) -> str:
"""Download model from HuggingFace Hub"""
print(f"Downloading model {model_id} from HuggingFace Hub")
# Use default cache dir if none provided
if cache_dir is None:
cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "helium", "models")
os.makedirs(cache_dir, exist_ok=True)
# Download model snapshot
local_path = snapshot_download(
repo_id=model_id,
cache_dir=cache_dir,
resume_download=True,
local_files_only=False # Always try to download first
)
print(f"Model downloaded to {local_path}")
return local_path
def store_in_device_db(
model_path: str,
device_db_url: str,
model_id: str
) -> None:
"""Store model weights and config in device DB"""
print(f"Storing model {model_id} in device DB at {device_db_url}")
# Connect to device DB
conn = duckdb.connect(device_db_url)
# Create tables if they don't exist
conn.execute("""
CREATE TABLE IF NOT EXISTS model_configs (
model_id VARCHAR PRIMARY KEY,
config JSON
);
CREATE TABLE IF NOT EXISTS model_weights (
model_id VARCHAR,
layer_name VARCHAR,
weight_data BLOB,
weight_dtype VARCHAR,
weight_shape VARCHAR,
PRIMARY KEY (model_id, layer_name)
);
""")
# Load and store config
config_path = os.path.join(model_path, "config.json")
with open(config_path, "r") as f:
config = json.load(f)
# Store config in DB
conn.execute(
"INSERT OR REPLACE INTO model_configs (model_id, config) VALUES (?, ?)",
[model_id, json.dumps(config)]
)
# Load and store weights
for root, _, files in os.walk(model_path):
for file in files:
if file.endswith(".bin"):
layer_name = os.path.splitext(file)[0]
weight_path = os.path.join(root, file)
# Load weight data
weight_data = np.fromfile(weight_path, dtype=np.float32)
weight_shape = weight_data.shape
# Store weight data in DB
conn.execute(
"""
INSERT OR REPLACE INTO model_weights
(model_id, layer_name, weight_data, weight_dtype, weight_shape)
VALUES (?, ?, ?, ?, ?)
""",
[model_id, layer_name, weight_data.tobytes(),
str(weight_data.dtype), str(weight_shape)]
)
conn.commit()
conn.execute(
"INSERT INTO model_configs (model_id, config) VALUES (?, ?)",
[model_id, json.dumps(config)]
)
# Load and store weights
weight_files = [f for f in os.listdir(model_path) if f.endswith(".bin")]
for weight_file in weight_files:
layer_name = os.path.splitext(weight_file)[0]
weight_path = os.path.join(model_path, weight_file)
# Load weights as numpy array
weights = np.load(weight_path)
# Store in DB
conn.execute(
"INSERT INTO model_weights (model_id, layer_name, weight_data) VALUES (?, ?, ?)",
[model_id, layer_name, weights.tobytes()]
)
conn.close()
|