File size: 2,897 Bytes
6f75615 | 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 | use pyo3::prelude::*;
use pyo3::exceptions::PyRuntimeError;
use candle_core::{Tensor, Device};
use std::collections::HashMap;
use std::fs::File;
use std::path::Path;
/// Helper function to perform raw tensor manipulation/quantization
fn perform_quantization(input_path: &str, output_path: &str, precision: &str) -> Result<(), String> {
let device = Device::Cpu; // Can switch to Device::new_cuda(0) for GPU
// 1. Read raw Llama 3.2 safetensors from the input directory
let content = std::fs::read(input_path)
.map_err(|e| format!("Failed to read input safetensor file: {}", e))?;
let safetensor = safetensors::SafeTensors::deserialize(&content)
.map_err(|e| format!("Invalid Safetensor format: {}", e))?;
let mut new_tensors = HashMap::new();
// 2. Iterate through layers (Weights, Attention blocks, Tensors)
for (name, tensor_view) in safetensor.tensors() {
let candle_tensor = candle_core::safetensors::load_buffer(tensor_view, &device)
.map_err(|e| format!("Candle failed to parse tensor {}: {}", name, e))?;
// 3. Low-level optimization based on UI input
let optimized_tensor = match precision {
"F16" => candle_tensor.to_dtype(candle_core::DType::F16)
.map_err(|e| e.to_string())?,
"BF16" => candle_tensor.to_dtype(candle_core::DType::BF16)
.map_err(|e| e.to_string())?,
_ => candle_tensor // Keep original if layout matches
};
new_tensors.insert(name, optimized_tensor);
}
// 4. Safely serialize and stream the weights back to disk
let output_file = Path::new(output_path);
candle_core::safetensors::save(&new_tensors, output_file)
.map_err(|e| format!("Failed to write optimized Llama 3.2 model: {}", e))?;
Ok(())
}
/// This is the actual function exposed directly to Python
#[pyfunction]
fn compress_safetensors(repo_id: String, precision: String) -> PyResult<bool> {
// Fake directory path routing for local example layout
let input_file = format!("./models/{}/model.safetensors", repo_id);
let output_file = format!("./models/{}/model_{}.safetensors", repo_id, precision.to_lowercase());
println!("[Rust Engine] Quantizing Llama 3.2 weights for: {}", repo_id);
match perform_quantization(&input_file, &output_file, &precision) {
Ok(_) => {
println!("[Rust Engine] Success. Optimized file written to {}", output_file);
Ok(true)
},
Err(err_msg) => {
// Converts native Rust system errors into explicit Python exceptions
Err(PyRuntimeError::new_err(err_msg))
}
}
}
/// Defines the Python library entry point structures
#[pymodule]
fn llama32_rust_core(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_function(wrap_pyfunction!(compress_safetensors, m)?)?;
Ok(())
}
|