hpt-dataloader 0.1.3

An internal library for data loading for hpt
Documentation 
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

use hpt_traits::tensor::{CommonBounds, TensorInfo};
use serde::{Deserialize, Serialize};
use std::fmt::Display;
use std::marker::PhantomData;
use std::{
    collections::HashMap,
    fs::File,
    io::{Read, Seek},
};

use crate::struct_save::load::load;
use crate::struct_save::load::MetaLoad;
use crate::struct_save::save::Save;
use crate::CPUTensorCreator;
use crate::CompressionAlgo;

#[derive(Debug, Serialize, Deserialize, Clone, Copy)]
pub enum Endian {
    Little,
    Big,
    Native,
}

#[derive(Debug, Serialize, Deserialize)]
pub(crate) struct HeaderInfos {
    pub(crate) begin: u64,
    pub(crate) infos: Vec<HeaderInfo>,
}

#[derive(Debug, Serialize, Deserialize)]
pub(crate) struct HeaderInfo {
    pub(crate) begin: u64,
    pub(crate) name: String,
    pub(crate) shape: Vec<i64>,
    pub(crate) strides: Vec<i64>,
    pub(crate) size: usize,
    pub(crate) indices: Vec<(usize, usize, usize, usize)>,
    pub(crate) compress_algo: CompressionAlgo,
    pub(crate) dtype: String,
    pub(crate) endian: Endian,
}
/// the meta data of the tensor
#[derive(Serialize, Deserialize)]
#[must_use]
pub struct TensorMeta<T: CommonBounds, B: CPUTensorCreator>
where
    <B as CPUTensorCreator>::Output: Clone + TensorInfo<T>,
{
    pub begin: usize,
    pub shape: Vec<i64>,
    pub strides: Vec<i64>,
    pub size: usize,
    pub dtype: String,
    pub compression_algo: CompressionAlgo,
    pub endian: Endian,
    pub indices: Vec<(usize, usize, usize, usize)>,
    pub phantom: PhantomData<(T, B)>,
}

pub fn parse_header_compressed<M: Save, P: Into<std::path::PathBuf>>(
    file: P,
) -> Result<<M as Save>::Meta, Box<dyn std::error::Error>> {
    let mut file = File::open(file.into())?;
    file.read_exact(&mut [0u8; "FASTTENSOR".len()])?;
    let mut header_infos = [0u8; 20];
    file.read_exact(&mut header_infos)?;
    let header = std::str::from_utf8(&header_infos)?;
    let header_int = header.trim().parse::<u64>()?;
    file.seek(std::io::SeekFrom::Start(header_int))?; // offset for header
    let mut buffer3 = vec![];
    file.read_to_end(&mut buffer3)?;
    let info = std::str::from_utf8(&buffer3)?;
    let ret = serde_json::from_str::<M::Meta>(info)?;
    Ok(ret)
}

impl<T, B: CPUTensorCreator> MetaLoad for TensorMeta<T, B>
where
    T: CommonBounds + bytemuck::AnyBitPattern,
    <B as CPUTensorCreator>::Output: Clone + TensorInfo<T> + Display + Into<B>,
{
    type Output = B;
    fn load(&self, file: &mut std::fs::File) -> std::io::Result<Self::Output> {
        Ok(load::<T, B>(file, self)?.into())
    }
}

impl HeaderInfo {
    pub(crate) fn parse_header_compressed(
        file: &str,
    ) -> Result<HashMap<String, HeaderInfo>, Box<dyn std::error::Error>> {
        let mut file = File::open(file)?;
        file.read_exact(&mut [0u8; "FASTTENSOR".len()])?;
        let mut header_infos = [0u8; 20];
        file.read_exact(&mut header_infos)?;
        let header = std::str::from_utf8(&header_infos)?;
        let header_int = header.trim().parse::<u64>()?;
        file.seek(std::io::SeekFrom::Start(header_int))?; // offset for header
        let mut buffer3 = vec![];
        file.read_to_end(&mut buffer3)?;
        let info = std::str::from_utf8(&buffer3)?;
        let ret: HashMap<String, HeaderInfo> =
            serde_json::from_str::<HashMap<String, Self>>(&info)?;
        Ok(ret)
    }
}