Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,675 @@
|
||||
//! 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(¤t_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(
|
||||
¤t_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
|
||||
));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user