huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-19 03:54:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

Tweak the kv-cache flag.

Browse Source
This commit is contained in:
laurent
2023-06-29 22:16:40 +01:00
parent ae3f202f3b
commit f6152e74b6
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
candle-core/examples/llama/main.rs
View File

@ -448,9 +448,9 @@ struct Args {
#[arg(long, default_value_t = 100)] #[arg(long, default_value_t = 100)]
sample_len: usize, sample_len: usize,
/// Enable the key-value cache. /// Disable the key-value cache.
#[arg(long, default_value_t = true)] #[arg(long)]
use_kv_cache: bool, no_kv_cache: bool,
} }
#[tokio::main] #[tokio::main]
@ -464,7 +464,7 @@ async fn main() -> Result<()> {
Device::new_cuda(0)? Device::new_cuda(0)?
}; };
let config = Config::config_7b(); let config = Config::config_7b();
let cache = Cache::new(args.use_kv_cache, &config, &device); let cache = Cache::new(!args.no_kv_cache, &config, &device);
let start = std::time::Instant::now(); let start = std::time::Instant::now();
let (llama, tokenizer_filename) = match args.npy { let (llama, tokenizer_filename) = match args.npy {
Some(npy) => { Some(npy) => {