janitor

janitor
git clone https://git.prime8.dev/janitor.git
Log | Files | Refs | README | LICENSE

tensor.rs (3038B)

      1 use std::cmp::Ordering;
      2 
      3 use anyhow::Result;
      4 use axum::http::StatusCode;
      5 use safetensors::{tensor::TensorView, Dtype};
      6 use tch::{Device, Kind, Tensor};
      7 
      8 pub const NUM_FRAMES: i64 = 1024;
      9 pub const NUM_MEL_BINS: i64 = 128;
     10 
     11 pub fn to_tensor(view: TensorView) -> Result<Tensor, (StatusCode, String)> {
     12     let size: Vec<i64> = view.shape().iter().map(|&x| x as i64).collect();
     13 
     14     let kind = match view.dtype() {
     15         Dtype::BOOL => Kind::Bool,
     16         Dtype::U8 => Kind::Uint8,
     17         Dtype::I8 => Kind::Int8,
     18         Dtype::I16 => Kind::Int16,
     19         Dtype::I32 => Kind::Int,
     20         Dtype::I64 => Kind::Int64,
     21         Dtype::BF16 => Kind::BFloat16,
     22         Dtype::F16 => Kind::Half,
     23         Dtype::F32 => Kind::Float,
     24         Dtype::F64 => Kind::Double,
     25         _ => {
     26             return Err((
     27                 StatusCode::BAD_REQUEST,
     28                 format!("unsupported dtype: {:?}", view.dtype()),
     29             ))
     30         }
     31     };
     32 
     33     let mut tensor = Tensor::f_from_data_size(view.data(), &size, kind)
     34         .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
     35 
     36     check(&tensor)?;
     37 
     38     tensor = tensor
     39         .f_to(Device::cuda_if_available())
     40         .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
     41 
     42     Ok(tensor)
     43 }
     44 
     45 pub fn check(tensor: &Tensor) -> Result<(), (StatusCode, String)> {
     46     let kind = tensor
     47         .f_kind()
     48         .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
     49 
     50     if kind != Kind::Float {
     51         return Err((
     52             StatusCode::BAD_REQUEST,
     53             "input must be a 32 bit float".to_string(),
     54         ));
     55     }
     56 
     57     let (_, num_mel_bins) = tensor
     58         .size2()
     59         .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
     60     if num_mel_bins != NUM_MEL_BINS {
     61         return Err((
     62             StatusCode::BAD_REQUEST,
     63             format!(
     64                 "Expected {} mel bins but got {}",
     65                 NUM_MEL_BINS, num_mel_bins
     66             ),
     67         ));
     68     }
     69     Ok(())
     70 }
     71 
     72 pub fn fit(tensor: &mut Tensor) -> Result<(), (StatusCode, String)> {
     73     let (num_frames, _) = tensor
     74         .size2()
     75         .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
     76 
     77     match NUM_FRAMES.cmp(&num_frames) {
     78         Ordering::Less => {
     79             *tensor = tensor
     80                 .f_narrow(0, 0, NUM_FRAMES)
     81                 .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
     82         }
     83         Ordering::Greater => {
     84             *tensor = tensor
     85                 .f_pad([0, 0, 0, NUM_FRAMES - num_frames], "constant", Some(0.0))
     86                 .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
     87         }
     88         _ => {}
     89     }
     90     Ok(())
     91 }
     92 
     93 pub fn normalize(tensor: &mut Tensor) -> Result<(), (StatusCode, String)> {
     94     *tensor = tensor
     95         .f_subtract_scalar(-4.2677393)
     96         .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
     97         .f_divide_scalar(4.5689974 * 2.0)
     98         .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
     99     Ok(())
    100 }