1import os
2import numpy as np
3from ._memobin import (
4 construct_dataset_url,
5 exists_in_memobin,
6 upload_to_memobin,
7)
10def upload_dataset_to_memobin(
11 data: np.ndarray,
12 dataset_name: str,
13 dataset_version: str,
14 memobin_api_key: str,
15 cache_dir: str,
16 verbose: bool = True,
17) -> None:
18 """Upload dataset to memobin in multiple formats.
20 Args:
21 data: The numpy array dataset to upload
22 dataset_name: Name of the dataset
23 dataset_version: Version of the dataset
24 memobin_api_key: API key for memobin
25 cache_dir: Directory for temporary files
26 verbose: Whether to print progress messages
27 """
28 try:
29 # Upload array metadata as JSON
30 dataset_url_json = construct_dataset_url(dataset_name, dataset_version, "json")
31 if not exists_in_memobin(dataset_url_json):
32 if verbose:
33 print(" Uploading dataset metadata to memobin...")
34 metadata = {"dtype": str(data.dtype), "shape": data.shape}
35 upload_to_memobin(
36 metadata,
37 dataset_url_json,
38 memobin_api_key,
39 content_type="application/json",
40 )
41 if verbose:
42 print(" Successfully uploaded metadata")
44 # Upload raw .dat format
45 dataset_url_raw = construct_dataset_url(dataset_name, dataset_version, "dat")
46 if not exists_in_memobin(dataset_url_raw):
47 if verbose:
48 print(" Uploading dataset (raw) to memobin...")
49 upload_to_memobin(
50 data.tobytes(),
51 dataset_url_raw,
52 memobin_api_key,
53 content_type="application/octet-stream",
54 )
55 if verbose:
56 print(" Successfully uploaded raw dataset")
58 # Upload .npy format
59 dataset_url_npy = construct_dataset_url(dataset_name, dataset_version, "npy")
60 if not exists_in_memobin(dataset_url_npy):
61 if verbose:
62 print(" Uploading dataset (npy) to memobin...")
63 # Save array to a temporary .npy file
64 temp_npy = os.path.join(cache_dir, "temp.npy")
65 np.save(temp_npy, data)
66 with open(temp_npy, "rb") as f:
67 npy_bytes = f.read()
68 os.remove(temp_npy) # Clean up temp file
70 upload_to_memobin(
71 npy_bytes,
72 dataset_url_npy,
73 memobin_api_key,
74 content_type="application/octet-stream",
75 )
76 if verbose:
77 print(" Successfully uploaded npy dataset")
78 except Exception as e:
79 print(f" Warning: Failed to upload dataset to memobin: {str(e)}")