gguf_rs_lib/builder/
tensor_builder.rs

1//! Tensor builder utilities
2
3use crate::error::{GGUFError, Result};
4use crate::tensor::{TensorData, TensorInfo, TensorShape, TensorType};
5
6#[cfg(all(not(feature = "std"), feature = "alloc"))]
7use hashbrown::HashSet;
8#[cfg(feature = "std")]
9use std::collections::HashSet;
10#[cfg(all(not(feature = "std"), feature = "alloc"))]
11extern crate alloc;
12#[cfg(all(not(feature = "std"), feature = "alloc"))]
13use alloc::{
14    format,
15    string::{String, ToString},
16    vec,
17    vec::Vec,
18};
19
20/// Builder for tensor collections
21#[cfg(any(feature = "std", feature = "alloc"))]
22#[derive(Debug, Default)]
23pub struct TensorCollectionBuilder {
24    tensors: Vec<(TensorInfo, TensorData)>,
25    names: HashSet<String>,
26}
27
28/// Builder for tensor collections (no_std + no_alloc variant)
29#[cfg(not(any(feature = "std", feature = "alloc")))]
30#[derive(Debug, Default)]
31pub struct TensorCollectionBuilder {
32    // Placeholder for no_std + no_alloc builds
33}
34
35#[cfg(any(feature = "std", feature = "alloc"))]
36impl TensorCollectionBuilder {
37    /// Create a new tensor collection builder
38    pub fn new() -> Self {
39        Self { tensors: Vec::new(), names: HashSet::new() }
40    }
41
42    /// Add a tensor
43    pub fn add_tensor<N: Into<String>>(
44        mut self,
45        name: N,
46        shape: Vec<u64>,
47        tensor_type: TensorType,
48        data: Vec<u8>,
49    ) -> Result<Self> {
50        let name = name.into();
51        if self.names.contains(&name) {
52            return Err(GGUFError::InvalidTensorData(format!("Duplicate tensor name: '{}'", name)));
53        }
54        let shape = TensorShape::new(shape)?;
55        let tensor_info = TensorInfo::new(name.clone(), shape, tensor_type, 0);
56        let tensor_data = TensorData::new_owned(data);
57
58        // Validate size
59        let expected_size =
60            usize::try_from(tensor_info.checked_expected_data_size()?).map_err(|_| {
61                GGUFError::InvalidTensorData("Tensor size does not fit this platform".to_string())
62            })?;
63        if tensor_data.len() != expected_size {
64            return Err(GGUFError::InvalidTensorData(format!(
65                "Size mismatch for tensor '{}'",
66                name
67            )));
68        }
69
70        tensor_info.validate()?;
71        tensor_data.validate()?;
72        self.tensors.try_reserve(1).map_err(|_| {
73            GGUFError::InvalidTensorData("Unable to allocate tensor collection".to_string())
74        })?;
75        self.names.try_reserve(1).map_err(|_| {
76            GGUFError::InvalidTensorData("Unable to allocate tensor-name index".to_string())
77        })?;
78        self.names.insert(name);
79        self.tensors.push((tensor_info, tensor_data));
80        Ok(self)
81    }
82
83    /// Add tensor with TensorData
84    pub fn add_tensor_data<N: Into<String>>(
85        mut self,
86        name: N,
87        shape: Vec<u64>,
88        tensor_type: TensorType,
89        data: TensorData,
90    ) -> Result<Self> {
91        let name = name.into();
92        if self.names.contains(&name) {
93            return Err(GGUFError::InvalidTensorData(format!("Duplicate tensor name: '{}'", name)));
94        }
95        let shape = TensorShape::new(shape)?;
96        let tensor_info = TensorInfo::new(name.clone(), shape, tensor_type, 0);
97
98        // Validate size
99        let expected_size =
100            usize::try_from(tensor_info.checked_expected_data_size()?).map_err(|_| {
101                GGUFError::InvalidTensorData("Tensor size does not fit this platform".to_string())
102            })?;
103        if data.len() != expected_size {
104            return Err(GGUFError::InvalidTensorData(format!(
105                "Size mismatch for tensor '{}'",
106                name
107            )));
108        }
109
110        tensor_info.validate()?;
111        data.validate()?;
112        self.tensors.try_reserve(1).map_err(|_| {
113            GGUFError::InvalidTensorData("Unable to allocate tensor collection".to_string())
114        })?;
115        self.names.try_reserve(1).map_err(|_| {
116            GGUFError::InvalidTensorData("Unable to allocate tensor-name index".to_string())
117        })?;
118        self.names.insert(name);
119        self.tensors.push((tensor_info, data));
120        Ok(self)
121    }
122
123    /// Build the tensor collection
124    pub fn build(self) -> Vec<(TensorInfo, TensorData)> {
125        self.tensors
126    }
127
128    /// Get tensor count
129    pub fn len(&self) -> usize {
130        self.tensors.len()
131    }
132
133    /// Check if empty
134    pub fn is_empty(&self) -> bool {
135        self.tensors.is_empty()
136    }
137
138    /// Check if tensor exists
139    pub fn contains(&self, name: &str) -> bool {
140        self.names.contains(name)
141    }
142}
143
144#[cfg(not(any(feature = "std", feature = "alloc")))]
145impl TensorCollectionBuilder {
146    /// Create a new tensor collection builder (no-op for no_std + no_alloc)
147    pub fn new() -> Self {
148        Self {}
149    }
150
151    /// Add a tensor (returns error for no_std + no_alloc)
152    pub fn add_tensor<N>(
153        self,
154        _name: N,
155        _shape: &[u64],
156        _tensor_type: TensorType,
157        _data: &[u8],
158    ) -> Result<Self> {
159        Err(GGUFError::AllocationRequired)
160    }
161
162    /// Add tensor with TensorData (returns error for no_std + no_alloc)
163    pub fn add_tensor_data<N>(
164        self,
165        _name: N,
166        _shape: &[u64],
167        _tensor_type: TensorType,
168        _data: &TensorData,
169    ) -> Result<Self> {
170        Err(GGUFError::AllocationRequired)
171    }
172
173    /// Build the tensor collection (returns empty for no_std + no_alloc)
174    pub fn build(self) -> &'static [(TensorInfo, TensorData)] {
175        &[]
176    }
177
178    /// Get tensor count (always 0 for no_std + no_alloc)
179    pub fn len(&self) -> usize {
180        0
181    }
182
183    /// Check if empty (always true for no_std + no_alloc)
184    pub fn is_empty(&self) -> bool {
185        true
186    }
187
188    /// Check if tensor exists (always false for no_std + no_alloc)
189    pub fn contains(&self, _name: &str) -> bool {
190        false
191    }
192}
193
194/// Helper for creating common tensor patterns
195pub struct TensorPatterns;
196
197impl TensorPatterns {
198    #[cfg(any(feature = "std", feature = "alloc"))]
199    fn validate_pair(
200        tensor_info: TensorInfo,
201        tensor_data: TensorData,
202    ) -> Result<(TensorInfo, TensorData)> {
203        tensor_info.validate()?;
204        tensor_data.validate()?;
205        let expected_size =
206            usize::try_from(tensor_info.checked_expected_data_size()?).map_err(|_| {
207                GGUFError::InvalidTensorData("Tensor size does not fit this platform".to_string())
208            })?;
209        if tensor_data.len() != expected_size {
210            return Err(GGUFError::InvalidTensorData(format!(
211                "Tensor '{}' data size mismatch: expected {}, got {}",
212                tensor_info.name(),
213                expected_size,
214                tensor_data.len()
215            )));
216        }
217        Ok((tensor_info, tensor_data))
218    }
219
220    /// Create a weight matrix tensor
221    #[cfg(any(feature = "std", feature = "alloc"))]
222    pub fn weight_matrix(
223        name: String,
224        input_dim: u64,
225        output_dim: u64,
226        tensor_type: TensorType,
227        data: Vec<u8>,
228    ) -> Result<(TensorInfo, TensorData)> {
229        let shape = TensorShape::new(vec![input_dim, output_dim])?;
230        let tensor_info = TensorInfo::new(name, shape, tensor_type, 0);
231        let tensor_data = TensorData::new_owned(data);
232        Self::validate_pair(tensor_info, tensor_data)
233    }
234
235    /// Create a bias vector tensor
236    #[cfg(any(feature = "std", feature = "alloc"))]
237    pub fn bias_vector(
238        name: String,
239        dim: u64,
240        tensor_type: TensorType,
241        data: Vec<u8>,
242    ) -> Result<(TensorInfo, TensorData)> {
243        let shape = TensorShape::new(vec![dim])?;
244        let tensor_info = TensorInfo::new(name, shape, tensor_type, 0);
245        let tensor_data = TensorData::new_owned(data);
246        Self::validate_pair(tensor_info, tensor_data)
247    }
248
249    /// Create an embedding matrix
250    #[cfg(any(feature = "std", feature = "alloc"))]
251    pub fn embedding_matrix(
252        name: String,
253        vocab_size: u64,
254        embedding_dim: u64,
255        tensor_type: TensorType,
256        data: Vec<u8>,
257    ) -> Result<(TensorInfo, TensorData)> {
258        Self::weight_matrix(name, embedding_dim, vocab_size, tensor_type, data)
259    }
260}
261
262#[cfg(all(test, feature = "std"))]
263mod tests {
264    use super::*;
265
266    #[test]
267    fn test_tensor_collection_builder() {
268        let collection = TensorCollectionBuilder::new()
269            .add_tensor("weight", vec![2, 3], TensorType::F32, vec![0u8; 24])
270            .unwrap()
271            .add_tensor("bias", vec![3], TensorType::F32, vec![0u8; 12])
272            .unwrap()
273            .build();
274
275        assert_eq!(collection.len(), 2);
276    }
277
278    #[test]
279    fn test_tensor_patterns() {
280        let (info, data) = TensorPatterns::weight_matrix(
281            "test_weight".to_string(),
282            4,
283            3,
284            TensorType::F32,
285            vec![0u8; 48], // 4*3*4 bytes
286        )
287        .unwrap();
288
289        assert_eq!(info.name(), "test_weight");
290        assert_eq!(info.shape().dims(), &[4, 3]);
291        assert_eq!(data.len(), 48);
292    }
293
294    #[test]
295    fn test_embedding_matrix_uses_ggml_dimension_order() {
296        let (info, _) = TensorPatterns::embedding_matrix(
297            "token_embd.weight".to_string(),
298            32_000,
299            2,
300            TensorType::F16,
301            vec![0; 32_000 * 2 * 2],
302        )
303        .unwrap();
304
305        assert_eq!(info.shape().dims(), &[2, 32_000]);
306    }
307
308    #[test]
309    fn test_tensor_collection_rejects_duplicates_and_preserves_order() {
310        let duplicate = TensorCollectionBuilder::new()
311            .add_tensor("first", vec![1], TensorType::F32, vec![0; 4])
312            .unwrap()
313            .add_tensor("first", vec![1], TensorType::F32, vec![0; 4]);
314        assert!(duplicate.is_err());
315
316        let collection = TensorCollectionBuilder::new()
317            .add_tensor("first", vec![1], TensorType::F32, vec![0; 4])
318            .unwrap()
319            .add_tensor("second", vec![1], TensorType::F32, vec![0; 4])
320            .unwrap()
321            .build();
322        assert_eq!(collection[0].0.name(), "first");
323        assert_eq!(collection[1].0.name(), "second");
324    }
325
326    #[test]
327    fn test_tensor_patterns_reject_wrong_payload_size() {
328        assert!(TensorPatterns::weight_matrix(
329            "bad".to_string(),
330            2,
331            3,
332            TensorType::F32,
333            vec![0; 4],
334        )
335        .is_err());
336        assert!(TensorPatterns::bias_vector("bad".to_string(), 3, TensorType::F32, vec![0; 4],)
337            .is_err());
338    }
339}
gguf_rs_lib/builder/tensor_builder.rs

gguf_rs_lib/builder/
tensor_builder.rs
