Add a verbose-prompt mode, similar to llama.cpp. (#489)

This commit is contained in:
Laurent Mazare
2023-08-17 15:26:44 +01:00
committed by GitHub
parent 5d99026fd2
commit ad7c53953b

View File

@ -303,6 +303,10 @@ struct Args {
/// Enable tracing (generates a trace-timestamp.json file). /// Enable tracing (generates a trace-timestamp.json file).
#[arg(long)] #[arg(long)]
tracing: bool, tracing: bool,
/// Display the token for the specified prompt.
#[arg(long)]
verbose_prompt: bool,
} }
impl Args { impl Args {
@ -375,11 +379,15 @@ fn main() -> anyhow::Result<()> {
let tokenizer = args.tokenizer()?; let tokenizer = args.tokenizer()?;
let prompt = args.prompt.as_ref().map_or(DEFAULT_PROMPT, |p| p.as_str()); let prompt = args.prompt.as_ref().map_or(DEFAULT_PROMPT, |p| p.as_str());
let mut tokens = tokenizer let tokens = tokenizer.encode(prompt, true).map_err(anyhow::Error::msg)?;
.encode(prompt, true) if args.verbose_prompt {
.map_err(anyhow::Error::msg)? for (token, id) in tokens.get_tokens().iter().zip(tokens.get_ids().iter()) {
.get_ids() let token = token.replace('▁', " ").replace("<0x0A>", "\n");
.to_vec(); println!("{id:7} -> '{token}'");
}
}
let mut tokens = tokens.get_ids().to_vec();
let mut index_pos = 0; let mut index_pos = 0;
let mut logits_processor = LogitsProcessor::new(args.seed, args.temperature); let mut logits_processor = LogitsProcessor::new(args.seed, args.temperature);
let start_gen = std::time::Instant::now(); let start_gen = std::time::Instant::now();