Load and run the T5-XXL encoder from a safetensors file. Returns hidden states (batch_size, seq_len, 4096). When `device` is a CUDA GPU with enough VRAM, runs T5 on GPU for ~5x speedup. Otherwise falls back to CPU with BF16 weights (~10GB RAM).
(
checkpoint_path: &std::path::Path,
prefix: &str,
input_ids: &Tensor,
device: &Device,
)
| 44 | /// When `device` is a CUDA GPU with enough VRAM, runs T5 on GPU for ~5x speedup. |
| 45 | /// Otherwise falls back to CPU with BF16 weights (~10GB RAM). |
| 46 | pub fn encode_t5( |
| 47 | checkpoint_path: &std::path::Path, |
| 48 | prefix: &str, |
| 49 | input_ids: &Tensor, |
| 50 | device: &Device, |
| 51 | ) -> Result<Tensor> { |
| 52 | let cfg = t5_xxl_config(); |
| 53 | |
| 54 | let run_on_gpu = matches!(device, Device::Cuda(_)); |
| 55 | |
| 56 | if run_on_gpu { |
| 57 | info!("loading T5-XXL text encoder (F16 on GPU)..."); |
| 58 | let vb = unsafe { |
| 59 | let filenames = vec![checkpoint_path.to_path_buf()]; |
| 60 | crate::utils::fp8::load_fp8_var_builder(&filenames, DType::F16, device)? |
| 61 | }; |
| 62 | let vb = vb.pp(prefix); |
| 63 | let input_ids = input_ids.to_device(device)?; |
| 64 | let mut model = T5EncoderModel::load(vb, &cfg)?; |
| 65 | info!("T5-XXL loaded on GPU, encoding..."); |
| 66 | let output = model.forward_dt(&input_ids, Some(DType::F32))?; |
| 67 | // Move output to CPU, then drop model to free GPU VRAM for transformer |
| 68 | let output = output.to_device(&Device::Cpu)?; |
| 69 | drop(model); |
| 70 | device.synchronize()?; |
| 71 | info!("T5-XXL encoding done, freed GPU memory"); |
| 72 | Ok(output) |
| 73 | } else { |
| 74 | info!("loading T5-XXL text encoder (BF16 on CPU, ~10GB RAM)..."); |
| 75 | let vb = unsafe { |
| 76 | let filenames = vec![checkpoint_path.to_path_buf()]; |
| 77 | crate::utils::fp8::load_fp8_var_builder(&filenames, DType::BF16, device)? |
| 78 | }; |
| 79 | let vb = vb.pp(prefix); |
| 80 | let mut model = T5EncoderModel::load(vb, &cfg)?; |
| 81 | info!("T5-XXL loaded, encoding..."); |
| 82 | let output = model.forward_dt(input_ids, Some(DType::F32))?; |
| 83 | Ok(output) |
| 84 | } |
| 85 | } |
no test coverage detected