debug: adjust channel to unbounded
parent
26a1a1e164
commit
019c745ac6
|
|
@ -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
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue