mirror of
https://github.com/huggingface/candle.git
synced 2025-06-15 10:26:33 +00:00
805bf9ffa7
* Implement top_p / nucleus sampling * Update changelog * rustfmt * Add tests * Fix clippy warning * Fix another clippy error
30 lines
939 B
Rust
30 lines
939 B
Rust
use candle::{Device, Result, Tensor};
|
|
use candle_transformers::generation::LogitsProcessor;
|
|
|
|
#[test]
|
|
fn sample_with_zero_temperature() -> Result<()> {
|
|
let mut logits_process = LogitsProcessor::new(1337, None, None);
|
|
let logits = Tensor::new(&[0.1, 0.2, 0.3, 0.4], &Device::Cpu)?;
|
|
let token = logits_process.sample(&logits)?;
|
|
assert_eq!(token, 3);
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn sample_with_temperature() -> Result<()> {
|
|
let mut logits_process = LogitsProcessor::new(42, Some(0.9), None);
|
|
let logits = Tensor::new(&[0.1, 0.2, 0.3, 0.4], &Device::Cpu)?;
|
|
let token = logits_process.sample(&logits)?;
|
|
assert_eq!(token, 0);
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn sample_with_top_p() -> Result<()> {
|
|
let mut logits_process = LogitsProcessor::new(42, Some(1.0), Some(0.5));
|
|
let logits = Tensor::new(&[0.1, 0.2, 0.3, 0.4], &Device::Cpu)?;
|
|
let token = logits_process.sample(&logits)?;
|
|
assert_eq!(token, 2);
|
|
Ok(())
|
|
}
|