| #define SAFETENSORS_CPP_IMPLEMENTATION |
|
|
| #include <cstdint> |
| #include <cstdlib> |
| #include <cstring> |
| #include <fstream> |
| #include <iostream> |
| #include <string> |
|
|
| #include "safetensors.hh" |
|
|
| int main(int argc, char **argv) { |
| if (argc != 2) { |
| return 2; |
| } |
|
|
| std::ifstream f(argv[1], std::ios::binary); |
| std::string contents((std::istreambuf_iterator<char>(f)), |
| std::istreambuf_iterator<char>()); |
| if (contents.empty()) { |
| return 2; |
| } |
|
|
| uint8_t *buf = static_cast<uint8_t *>(std::malloc(contents.size())); |
| if (!buf) { |
| return 2; |
| } |
| std::memcpy(buf, contents.data(), contents.size()); |
|
|
| safetensors::safetensors_t st; |
| std::string warn; |
| std::string err; |
| bool ok = safetensors::mmap_from_memory(buf, contents.size(), argv[1], &st, |
| &warn, &err); |
| std::cerr << "mmap_from_memory=" << ok << " err=" << err << "\n"; |
| if (!ok) { |
| return 1; |
| } |
|
|
| err.clear(); |
| bool valid = safetensors::validate_data_offsets(st, err); |
| std::cerr << "validate_data_offsets=" << valid << " err=" << err << "\n"; |
|
|
| safetensors::tensor_t a; |
| safetensors::tensor_t b; |
| if (!st.tensors.at("a", &a) || !st.tensors.at("b", &b)) { |
| return 1; |
| } |
|
|
| const uint8_t *base = st.databuffer_addr; |
| std::cerr << "databuffer_size=" << st.databuffer_size << "\n"; |
| std::cerr << "a=[" << a.data_offsets[0] << "," << a.data_offsets[1] |
| << "] bytes=" << char(base[a.data_offsets[0]]) |
| << char(base[a.data_offsets[1] - 1]) << "\n"; |
| std::cerr << "b=[" << b.data_offsets[0] << "," << b.data_offsets[1] |
| << "] byte=" << char(base[b.data_offsets[0]]) << "\n"; |
| std::cerr << "alias=" << (base + b.data_offsets[0] == base + a.data_offsets[1] - 1) |
| << "\n"; |
|
|
| std::free(buf); |
| return valid ? 0 : 1; |
| } |
|
|
