huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-17 11:08:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix build issue in EOS Token in llama-multiprocess (#2420)

Browse Source
This commit is contained in:
Hadi
2024-08-16 12:46:31 -04:00
committed by GitHub
parent 53ce65f706
commit 2b75dd9551
1 changed files with 11 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
candle-examples/examples/llama_multiprocess/main.rs
View File

@ -14,12 +14,14 @@ use clap::{Parser, ValueEnum};
use candle::{DType, Device, Tensor}; use candle::{DType, Device, Tensor};
use candle_transformers::generation::LogitsProcessor; use candle_transformers::generation::LogitsProcessor;
use candle_transformers::models::llama::LlamaEosToks;
use cudarc::driver::safe::CudaDevice; use cudarc::driver::safe::CudaDevice;
use cudarc::nccl::safe::{Comm, Id}; use cudarc::nccl::safe::{Comm, Id};
use hf_hub::{api::sync::Api, Repo, RepoType}; use hf_hub::{api::sync::Api, Repo, RepoType};
use std::io::Write; use std::io::Write;
use std::rc::Rc; use std::rc::Rc;
mod model; mod model;
use model::{Config, Llama}; use model::{Config, Llama};
@ -219,9 +221,16 @@ fn main() -> Result<()> {
let next_token = logits_processor.sample(&logits)?; let next_token = logits_processor.sample(&logits)?;
tokens.push(next_token); tokens.push(next_token);
new_tokens.push(next_token); new_tokens.push(next_token);
if Some(next_token) == config.eos_token_id { match config.eos_token_id {
Some(LlamaEosToks::Single(eos_tok_id)) if next_token == eos_tok_id => {
break; break;
} }
Some(LlamaEosToks::Multiple(ref eos_ids)) if eos_ids.contains(&next_token) => {
break;
}
_ => (),
}
if rank == 0 { if rank == 0 {
if let Some(t) = tokenizer.next_token(next_token)? { if let Some(t) = tokenizer.next_token(next_token)? {
print!("{t}"); print!("{t}");