debug: adjust channel to unbounded

Meng Zhang 2023-10-09 18:43:26 -07:00
parent 26a1a1e164
commit 019c745ac6
1 changed files with 5 additions and 5 deletions

View File

@ -9,7 +9,7 @@ use tabby_inference::{
helpers, TextGeneration, TextGenerationOptions, helpers, TextGeneration, TextGenerationOptions,
}; };
use tokenizers::tokenizer::Tokenizer; use tokenizers::tokenizer::Tokenizer;
use tokio::sync::mpsc::{channel, Sender}; use tokio::sync::mpsc::{unbounded_channel, UnboundedSender};
use tokio_util::sync::CancellationToken; use tokio_util::sync::CancellationToken;
#[cxx::bridge(namespace = "tabby")] #[cxx::bridge(namespace = "tabby")]
@ -72,14 +72,14 @@ pub struct CTranslate2EngineOptions {
} }
pub struct InferenceContext { pub struct InferenceContext {
sender: Sender<String>, sender: UnboundedSender<String>,
decoding: IncrementalDecoding, decoding: IncrementalDecoding,
cancel: CancellationToken, cancel: CancellationToken,
} }
impl InferenceContext { impl InferenceContext {
fn new( fn new(
sender: Sender<String>, sender: UnboundedSender<String>,
decoding: IncrementalDecoding, decoding: IncrementalDecoding,
cancel: CancellationToken, cancel: CancellationToken,
) -> Self { ) -> Self {
@ -137,7 +137,7 @@ impl TextGeneration for CTranslate2Engine {
let cancel = CancellationToken::new(); let cancel = CancellationToken::new();
let engine = self.engine.clone(); let engine = self.engine.clone();
let (sender, mut receiver) = channel::<String>(8); let (sender, mut receiver) = unbounded_channel();
let context = InferenceContext::new(sender, decoding, cancel.clone()); let context = InferenceContext::new(sender, decoding, cancel.clone());
tokio::task::spawn(async move { tokio::task::spawn(async move {
let context = Box::new(context); let context = Box::new(context);
@ -178,7 +178,7 @@ fn inference_callback(
if context.cancel.is_cancelled() { if context.cancel.is_cancelled() {
true true
} else if let Some(new_text) = context.decoding.next_token(token_id) { } else if let Some(new_text) = context.decoding.next_token(token_id) {
let _ = context.sender.blocking_send(new_text); let _ = context.sender.send(new_text);
false false
} else { } else {
true true