Files
rustytorch/crates/models/rtx-nlg/src/serving/streaming.rs
T
2026-03-04 00:08:42 +00:00

676 lines
22 KiB
Rust

//! Streaming text generation with token-by-token output and real-time processing
use crate::{GenerationConfig, GenerationOutput, ModelInterface, Result, generation::utils};
use futures::Stream;
use rtx_tensor::{Device, Tensor};
use std::collections::VecDeque;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::mpsc;
/// Streaming text generator for real-time token-by-token generation
pub struct StreamingGenerator {
model: Arc<dyn ModelInterface>,
config: GenerationConfig,
streaming_config: StreamingConfig,
}
/// Configuration for streaming generation
#[derive(Debug, Clone)]
pub struct StreamingConfig {
/// Buffer size for streaming tokens
pub buffer_size: usize,
/// Maximum latency per token in milliseconds
pub max_token_latency_ms: u64,
/// Whether to stream partial tokens (subword pieces)
pub stream_partial_tokens: bool,
/// Minimum token confidence to stream
pub min_confidence: f32,
/// Whether to include timing information
pub include_timing: bool,
/// Whether to include confidence scores
pub include_confidence: bool,
/// Batch size for streaming (1 for pure token-by-token)
pub batch_size: usize,
/// Whether to use speculative decoding for faster streaming
pub use_speculative_decoding: bool,
}
impl Default for StreamingConfig {
fn default() -> Self {
Self {
buffer_size: 100,
max_token_latency_ms: 100,
stream_partial_tokens: false,
min_confidence: 0.0,
include_timing: true,
include_confidence: false,
batch_size: 1,
use_speculative_decoding: false,
}
}
}
impl StreamingGenerator {
/// Create streaming generator with model and generation config
pub fn with_model(model: Arc<dyn ModelInterface>, config: GenerationConfig) -> Result<Self> {
Self::with_streaming_config(model, config, StreamingConfig::default())
}
/// Create streaming generator with custom streaming configuration
pub fn with_streaming_config(
model: Arc<dyn ModelInterface>,
config: GenerationConfig,
streaming_config: StreamingConfig,
) -> Result<Self> {
config.validate()?;
Ok(Self {
model,
config,
streaming_config,
})
}
/// Generate streaming text with real-time token output
pub async fn generate_stream(
&self,
prompt: &str,
) -> Result<impl Stream<Item = Result<StreamingToken>>> {
let (tx, rx) = mpsc::channel(self.streaming_config.buffer_size);
let model = self.model.clone();
let config = self.config.clone();
let streaming_config = self.streaming_config.clone();
let prompt = prompt.to_string();
// Spawn background task for generation
tokio::spawn(async move {
if let Err(e) =
Self::generate_streaming_impl(model, &prompt, config, streaming_config, tx).await
{
tracing::error!("Streaming generation failed: {}", e);
}
});
Ok(tokio_stream::wrappers::ReceiverStream::new(rx))
}
/// Generate multiple streaming completions
pub async fn generate_multiple_streams(
&self,
prompt: &str,
num_streams: usize,
) -> Result<Vec<tokio_stream::wrappers::ReceiverStream<Result<StreamingToken>>>> {
let mut streams = Vec::new();
let prompt_owned = prompt.to_string();
for i in 0..num_streams {
let seed = self.config.seed.map(|s| s + i as u64);
let mut modified_config = self.config.clone();
modified_config.seed = seed;
// Clone everything needed for the spawned task
let model = self.model.clone();
let streaming_config = self.streaming_config.clone();
let prompt_clone = prompt_owned.clone();
let (tx, rx) = mpsc::channel(self.streaming_config.buffer_size);
tokio::spawn(async move {
if let Err(e) = Self::generate_streaming_impl(
model,
&prompt_clone,
modified_config,
streaming_config,
tx,
)
.await
{
tracing::error!("Streaming generation failed: {}", e);
}
});
streams.push(tokio_stream::wrappers::ReceiverStream::new(rx));
}
Ok(streams)
}
/// Internal streaming implementation
async fn generate_streaming_impl(
model: Arc<dyn ModelInterface>,
prompt: &str,
config: GenerationConfig,
streaming_config: StreamingConfig,
tx: mpsc::Sender<Result<StreamingToken>>,
) -> Result<()> {
let start_time = Instant::now();
let tokenizer = model.tokenizer();
let special_tokens = tokenizer.special_tokens();
// Tokenize input
let input_tokens = tokenizer.encode(prompt)?;
let mut current_sequence = input_tokens.clone();
let max_length = config.max_length.unwrap_or(100);
let max_new_tokens = max_length.saturating_sub(input_tokens.len());
// Send initial metadata token
let metadata_token = StreamingToken {
text: String::new(),
token_id: 0,
position: 0,
confidence: 1.0,
timing: TokenTiming {
generation_time_ms: 0.0,
cumulative_time_ms: 0.0,
tokens_per_second: 0.0,
},
metadata: Some(StreamingMetadata {
prompt_tokens: input_tokens.len(),
max_new_tokens,
is_start: true,
is_end: false,
finish_reason: None,
}),
token_type: TokenType::Metadata,
};
if let Err(_) = tx.send(Ok(metadata_token)).await {
return Ok(()); // Receiver dropped
}
// Main generation loop
let mut generated_tokens = 0;
let mut last_token_time = start_time;
while generated_tokens < max_new_tokens {
let token_start = Instant::now();
// Create input tensor for current sequence
let sequence_data: Vec<f32> = current_sequence.iter().map(|&x| x as f32).collect();
let input_tensor = Tensor::from_data(
sequence_data,
[1, current_sequence.len()],
&Device::default(),
)?;
// Get model predictions (single forward pass)
let model_output = if streaming_config.use_speculative_decoding {
// Use speculative decoding for faster generation
generate_with_speculative_decoding(&*model, &input_tensor, &config).await?
} else {
model.generate_tokens(&input_tensor, None, &config)?
};
// Get logits for next token
let logits = get_next_token_logits(&model_output, &input_tensor)?;
// Apply generation strategy (temperature, top-p, top-k, etc.)
let (next_token, confidence) =
sample_next_token_with_confidence(&logits, &config, &streaming_config)?;
// Check if we should stop generation
if utils::should_stop_generation(&current_sequence, &config, special_tokens) {
break;
}
// Check confidence threshold
if confidence < streaming_config.min_confidence {
break;
}
// Add token to sequence
current_sequence.push(next_token);
generated_tokens += 1;
// Decode token text
let token_text = if streaming_config.stream_partial_tokens {
// Stream individual token pieces
tokenizer.decode(&[next_token])?
} else {
// Stream full token only
let decoded = tokenizer.decode(&[next_token])?;
if is_complete_token(&decoded) {
decoded
} else {
String::new() // Skip partial tokens
}
};
let token_end = Instant::now();
let token_time = token_end.duration_since(token_start);
let cumulative_time = token_end.duration_since(start_time);
// Calculate timing metrics
let timing = TokenTiming {
generation_time_ms: token_time.as_millis() as f64,
cumulative_time_ms: cumulative_time.as_millis() as f64,
tokens_per_second: generated_tokens as f64 / cumulative_time.as_secs_f64(),
};
// Create streaming token
let streaming_token = StreamingToken {
text: token_text,
token_id: next_token,
position: input_tokens.len() + generated_tokens - 1,
confidence: if streaming_config.include_confidence {
confidence
} else {
1.0
},
timing: if streaming_config.include_timing {
timing
} else {
TokenTiming::default()
},
metadata: None,
token_type: TokenType::Generated,
};
// Send token if not empty or if it's a special token
if (!streaming_token.text.is_empty() || is_special_token(next_token, special_tokens))
&& let Err(_) = tx.send(Ok(streaming_token)).await
{
break; // Receiver dropped
}
// Rate limiting - ensure we don't exceed max token latency
let elapsed_since_last = token_end.duration_since(last_token_time);
let min_interval = Duration::from_millis(streaming_config.max_token_latency_ms);
if elapsed_since_last < min_interval {
tokio::time::sleep(min_interval.checked_sub(elapsed_since_last).unwrap()).await;
}
last_token_time = Instant::now();
// Check for early stopping conditions
if next_token == special_tokens.eos_token {
break;
}
// Check timeout
if let Some(timeout_ms) = config.timeout_ms
&& cumulative_time.as_millis() > timeout_ms as u128
{
break;
}
}
// Send final metadata token
let final_time = Instant::now().duration_since(start_time);
let final_metadata = StreamingToken {
text: String::new(),
token_id: 0,
position: current_sequence.len(),
confidence: 1.0,
timing: TokenTiming {
generation_time_ms: final_time.as_millis() as f64,
cumulative_time_ms: final_time.as_millis() as f64,
tokens_per_second: generated_tokens as f64 / final_time.as_secs_f64(),
},
metadata: Some(StreamingMetadata {
prompt_tokens: input_tokens.len(),
max_new_tokens,
is_start: false,
is_end: true,
finish_reason: Some(determine_finish_reason(
&current_sequence,
&config,
special_tokens,
)),
}),
token_type: TokenType::Metadata,
};
let _ = tx.send(Ok(final_metadata)).await;
Ok(())
}
}
/// Streaming token with timing and confidence information
#[derive(Debug, Clone)]
pub struct StreamingToken {
/// Generated text for this token
pub text: String,
/// Token ID
pub token_id: u32,
/// Position in the sequence
pub position: usize,
/// Confidence score for this token
pub confidence: f32,
/// Timing information
pub timing: TokenTiming,
/// Metadata (for start/end tokens)
pub metadata: Option<StreamingMetadata>,
/// Type of token
pub token_type: TokenType,
}
/// Timing information for each token
#[derive(Debug, Clone, Default)]
pub struct TokenTiming {
/// Time to generate this token in milliseconds
pub generation_time_ms: f64,
/// Cumulative generation time in milliseconds
pub cumulative_time_ms: f64,
/// Current generation speed in tokens per second
pub tokens_per_second: f64,
}
/// Metadata for streaming session
#[derive(Debug, Clone)]
pub struct StreamingMetadata {
/// Number of prompt tokens
pub prompt_tokens: usize,
/// Maximum new tokens to generate
pub max_new_tokens: usize,
/// Whether this is the start of generation
pub is_start: bool,
/// Whether this is the end of generation
pub is_end: bool,
/// Reason for finishing generation
pub finish_reason: Option<crate::FinishReason>,
}
/// Type of streaming token
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TokenType {
/// Regular generated token
Generated,
/// Metadata token (start/end)
Metadata,
/// Special token (EOS, etc.)
Special,
}
/// Buffer for accumulating partial tokens
#[derive(Debug)]
pub struct TokenBuffer {
buffer: VecDeque<String>,
max_size: usize,
}
impl TokenBuffer {
pub fn new(max_size: usize) -> Self {
Self {
buffer: VecDeque::with_capacity(max_size),
max_size,
}
}
pub fn add_token(&mut self, token: String) {
if self.buffer.len() >= self.max_size {
self.buffer.pop_front();
}
self.buffer.push_back(token);
}
pub fn get_text(&self) -> String {
self.buffer.iter().cloned().collect::<String>()
}
pub fn clear(&mut self) {
self.buffer.clear();
}
}
// Helper functions
fn get_next_token_logits(_output: &GenerationOutput, _input: &Tensor) -> Result<Tensor> {
// Mock implementation - in real scenario would extract logits from model output
let vocab_size = 50000;
// Create streaming-optimized logits with quick token selection characteristics
let mut logits_vec = Vec::with_capacity(vocab_size);
for i in 0..vocab_size {
let base_logit = rand::random::<f32>() * 2.5 - 1.25; // Random between -1.25 and 1.25
// Streaming favors more common tokens for faster generation
let frequency_boost = if i < 2000 {
1.0
} else if i < 10000 {
0.3
} else {
-0.5
};
logits_vec.push(base_logit + frequency_boost);
}
let logits = Tensor::from_data(logits_vec, [vocab_size], &Device::default())?;
Ok(logits)
}
fn sample_next_token_with_confidence(
logits: &Tensor,
config: &GenerationConfig,
_streaming_config: &StreamingConfig,
) -> Result<(u32, f32)> {
let mut modified_logits = logits.clone();
// Apply temperature
utils::apply_temperature(&mut modified_logits, config.temperature)?;
// Apply sampling based on generation strategy
match &config.strategy {
crate::generation::GenerationStrategy::NucleusSampling(nucleus_config) => {
utils::apply_top_p(&mut modified_logits, nucleus_config.top_p)?;
}
crate::generation::GenerationStrategy::TopKSampling(topk_config) => {
utils::apply_top_k(&mut modified_logits, topk_config.top_k)?;
}
_ => {} // Use raw probabilities for other strategies
}
// Sample token and calculate confidence
let mut rng = rand::thread_rng();
let next_token = utils::sample_categorical(&modified_logits, &mut rng)?;
// Calculate confidence as the probability of the selected token
let probs = modified_logits.softmax(-1)?;
let confidence = probs.get(&[next_token as usize])?;
Ok((next_token, confidence))
}
async fn generate_with_speculative_decoding(
model: &dyn ModelInterface,
input: &Tensor,
config: &GenerationConfig,
) -> Result<GenerationOutput> {
// Simplified speculative decoding - would implement draft model + verification
model.generate_tokens(input, None, config)
}
fn is_complete_token(text: &str) -> bool {
// Heuristic to determine if a token is complete (not a subword piece)
!text.starts_with("##") && !text.is_empty()
}
fn is_special_token(token_id: u32, special_tokens: &crate::SpecialTokens) -> bool {
token_id == special_tokens.eos_token
|| token_id == special_tokens.bos_token
|| token_id == special_tokens.pad_token
|| token_id == special_tokens.unk_token
|| (special_tokens.sep_token == Some(token_id))
|| (special_tokens.cls_token == Some(token_id))
|| (special_tokens.mask_token == Some(token_id))
}
fn determine_finish_reason(
sequence: &[u32],
config: &GenerationConfig,
special_tokens: &crate::SpecialTokens,
) -> crate::FinishReason {
if let Some(&last_token) = sequence.last()
&& last_token == special_tokens.eos_token
{
return crate::FinishReason::EosToken;
}
if let Some(max_length) = config.max_length
&& sequence.len() >= max_length
{
return crate::FinishReason::MaxLength;
}
crate::FinishReason::MaxLength
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{GenerationConfig, MockModelInterface};
use futures::StreamExt;
#[tokio::test]
async fn test_streaming_generation() -> Result<()> {
let model = Arc::new(MockModelInterface::new("/tmp/mock")?);
let config = GenerationConfig::default().max_length(10);
let generator = StreamingGenerator::with_model(model, config)?;
let mut stream = generator.generate_stream("test prompt").await?;
let mut tokens = Vec::new();
while let Some(result) = stream.next().await {
tokens.push(result?);
}
assert!(!tokens.is_empty());
// Check for metadata tokens
let start_metadata = tokens.iter().find(|t| {
matches!(t.token_type, TokenType::Metadata)
&& t.metadata.as_ref().map_or(false, |m| m.is_start)
});
let end_metadata = tokens.iter().find(|t| {
matches!(t.token_type, TokenType::Metadata)
&& t.metadata.as_ref().map_or(false, |m| m.is_end)
});
assert!(start_metadata.is_some());
assert!(end_metadata.is_some());
Ok(())
}
#[tokio::test]
async fn test_streaming_with_custom_config() -> Result<()> {
let model = Arc::new(MockModelInterface::new("/tmp/mock")?);
let config = GenerationConfig::default().max_length(5);
let streaming_config = StreamingConfig {
buffer_size: 50,
max_token_latency_ms: 10,
include_timing: true,
include_confidence: true,
..Default::default()
};
let generator = StreamingGenerator::with_streaming_config(model, config, streaming_config)?;
let mut stream = generator.generate_stream("hello").await?;
let mut tokens = Vec::new();
while let Some(result) = stream.next().await {
tokens.push(result?);
}
assert!(!tokens.is_empty());
// Check that timing and confidence are included
let generated_tokens: Vec<_> = tokens
.iter()
.filter(|t| matches!(t.token_type, TokenType::Generated))
.collect();
for token in generated_tokens {
assert!(token.timing.generation_time_ms >= 0.0);
assert!(token.confidence >= 0.0);
}
Ok(())
}
#[tokio::test]
async fn test_multiple_streams() -> Result<()> {
let model = Arc::new(MockModelInterface::new("/tmp/mock")?);
let config = GenerationConfig::default().max_length(5).seed(123);
let generator = StreamingGenerator::with_model(model, config)?;
let streams = generator.generate_multiple_streams("test", 2).await?;
assert_eq!(streams.len(), 2);
Ok(())
}
#[test]
fn test_token_buffer() {
let mut buffer = TokenBuffer::new(3);
buffer.add_token("hello".to_string());
buffer.add_token(" ".to_string());
buffer.add_token("world".to_string());
assert_eq!(buffer.get_text(), "hello world");
buffer.add_token("!".to_string());
assert_eq!(buffer.get_text(), " world!"); // First token evicted
buffer.clear();
assert_eq!(buffer.get_text(), "");
}
#[test]
fn test_is_complete_token() {
assert!(is_complete_token("hello"));
assert!(is_complete_token("world"));
assert!(!is_complete_token("##ing")); // Subword piece
assert!(!is_complete_token("")); // Empty
}
#[test]
fn test_is_special_token() {
let special_tokens = crate::SpecialTokens {
pad_token: 0,
eos_token: 2,
bos_token: 1,
unk_token: 3,
sep_token: Some(4),
cls_token: Some(5),
mask_token: Some(6),
};
assert!(is_special_token(0, &special_tokens)); // PAD
assert!(is_special_token(2, &special_tokens)); // EOS
assert!(is_special_token(4, &special_tokens)); // SEP
assert!(!is_special_token(100, &special_tokens)); // Regular token
}
#[test]
fn test_determine_finish_reason() {
let special_tokens = crate::SpecialTokens {
pad_token: 0,
eos_token: 2,
bos_token: 1,
unk_token: 3,
sep_token: None,
cls_token: None,
mask_token: None,
};
let config = GenerationConfig::default().max_length(5);
// EOS token finish
let eos_sequence = vec![1, 10, 20, 2];
assert!(matches!(
determine_finish_reason(&eos_sequence, &config, &special_tokens),
crate::FinishReason::EosToken
));
// Max length finish
let long_sequence = vec![1, 10, 20, 30, 40];
assert!(matches!(
determine_finish_reason(&long_sequence, &config, &special_tokens),
crate::FinishReason::MaxLength
));
}
}