384 lines
13 KiB
Rust
384 lines
13 KiB
Rust
//! Integration benchmarks for advanced NLG features
|
|
|
|
use criterion::{BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main};
|
|
use futures::StreamExt;
|
|
use rtx_nlg::*;
|
|
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
|
|
fn bench_streaming_generation(c: &mut Criterion) {
|
|
let rt = tokio::runtime::Runtime::new().unwrap();
|
|
|
|
let mut group = c.benchmark_group("streaming_generation");
|
|
group.throughput(Throughput::Elements(1));
|
|
|
|
for max_length in [25, 50, 100].iter() {
|
|
group.bench_with_input(
|
|
BenchmarkId::new("max_length", max_length),
|
|
max_length,
|
|
|b, &max_length| {
|
|
b.to_async(&rt).iter(|| async {
|
|
let model = Arc::new(MockModelInterface::new("/tmp/mock").unwrap());
|
|
let config = GenerationConfig::default().max_length(max_length);
|
|
let generator = serving::StreamingGenerator::with_model(model, config).unwrap();
|
|
|
|
let mut stream = generator
|
|
.generate_stream(black_box("Tell me about AI"))
|
|
.await
|
|
.unwrap();
|
|
let mut tokens = Vec::new();
|
|
|
|
while let Some(token) = stream.next().await {
|
|
tokens.push(token.unwrap());
|
|
}
|
|
|
|
black_box(tokens);
|
|
});
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_constrained_generation(c: &mut Criterion) {
|
|
let rt = tokio::runtime::Runtime::new().unwrap();
|
|
|
|
let mut group = c.benchmark_group("constrained_generation");
|
|
group.throughput(Throughput::Elements(1));
|
|
|
|
let lexical_constraint = generation::constrained::LexicalConstraint::new()
|
|
.add_required_phrase(vec![100, 200, 300]) // Mock token sequence
|
|
.add_forbidden_phrase(vec![400, 500]);
|
|
|
|
let constrained_config = generation::constrained::ConstrainedConfig {
|
|
lexical_constraints: vec![lexical_constraint],
|
|
satisfaction_strategy: generation::constrained::ConstraintStrategy::Soft,
|
|
..Default::default()
|
|
};
|
|
|
|
group.bench_function("with_constraints", |b| {
|
|
b.to_async(&rt).iter(|| async {
|
|
let model = Arc::new(MockModelInterface::new("/tmp/mock").unwrap());
|
|
let config = GenerationConfig::default().max_length(50);
|
|
let input_tensor = rtx_tensor::Tensor::from_slice(&[1u32, 2, 3], &[1, 3]).unwrap();
|
|
|
|
let result = generation::constrained::generate_constrained(
|
|
&*model,
|
|
&input_tensor,
|
|
&config,
|
|
&constrained_config,
|
|
)
|
|
.unwrap();
|
|
|
|
black_box(result);
|
|
});
|
|
});
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_controllable_generation(c: &mut Criterion) {
|
|
let rt = tokio::runtime::Runtime::new().unwrap();
|
|
|
|
let mut group = c.benchmark_group("controllable_generation");
|
|
group.throughput(Throughput::Elements(1));
|
|
|
|
let sentiment_control = generation::controllable::SentimentControl::new(0.8)
|
|
.add_positive_word(100, 1.0)
|
|
.add_positive_word(200, 0.8);
|
|
|
|
let controllable_config = generation::controllable::ControllableConfig {
|
|
sentiment_control: Some(sentiment_control),
|
|
control_strength: 0.7,
|
|
..Default::default()
|
|
};
|
|
|
|
for control_method in [
|
|
generation::controllable::ControlMethod::WeightedDecoding,
|
|
generation::controllable::ControlMethod::ClassifierGuided,
|
|
]
|
|
.iter()
|
|
{
|
|
let mut config_with_method = controllable_config.clone();
|
|
config_with_method.control_method = control_method.clone();
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("method", format!("{:?}", control_method)),
|
|
&config_with_method,
|
|
|b, config| {
|
|
b.to_async(&rt).iter(|| async {
|
|
let model = Arc::new(MockModelInterface::new("/tmp/mock").unwrap());
|
|
let gen_config = GenerationConfig::default().max_length(50);
|
|
let input_tensor =
|
|
rtx_tensor::Tensor::from_slice(&[1u32, 2, 3], &[1, 3]).unwrap();
|
|
|
|
let result = generation::controllable::generate_controllable(
|
|
&*model,
|
|
&input_tensor,
|
|
&gen_config,
|
|
&config,
|
|
)
|
|
.unwrap();
|
|
|
|
black_box(result);
|
|
});
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_template_rendering(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("template_rendering");
|
|
group.throughput(Throughput::Elements(1));
|
|
|
|
let mut manager = serving::templates::TemplateManager::new();
|
|
|
|
// Add complex template
|
|
let few_shot_examples = vec![
|
|
serving::templates::FewShotExample {
|
|
input: "positive example".to_string(),
|
|
output: "positive".to_string(),
|
|
explanation: Some("This is positive".to_string()),
|
|
},
|
|
serving::templates::FewShotExample {
|
|
input: "negative example".to_string(),
|
|
output: "negative".to_string(),
|
|
explanation: Some("This is negative".to_string()),
|
|
},
|
|
];
|
|
|
|
let complex_template = serving::templates::PromptTemplate::few_shot(
|
|
"complex_classification".to_string(),
|
|
"Classify: {{text}}".to_string(),
|
|
few_shot_examples,
|
|
);
|
|
|
|
manager.add_template(complex_template);
|
|
|
|
let test_values = [
|
|
("basic", vec![("name", "Alice"), ("age", "30")]),
|
|
(
|
|
"medium",
|
|
vec![
|
|
("text", "This is a test sentence"),
|
|
("context", "Testing context"),
|
|
],
|
|
),
|
|
(
|
|
"complex",
|
|
vec![("text", "Complex classification task"), ("domain", "AI")],
|
|
),
|
|
];
|
|
|
|
for (complexity, value_pairs) in test_values.iter() {
|
|
let mut values = HashMap::new();
|
|
for (k, v) in value_pairs {
|
|
values.insert(k.to_string(), v.to_string());
|
|
}
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("complexity", complexity),
|
|
&values,
|
|
|b, values| {
|
|
b.iter(|| {
|
|
let result = manager
|
|
.render_template(black_box("complex_classification"), black_box(values))
|
|
.unwrap();
|
|
black_box(result);
|
|
});
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_parallel_generation(c: &mut Criterion) {
|
|
let rt = tokio::runtime::Runtime::new().unwrap();
|
|
|
|
let mut group = c.benchmark_group("parallel_generation");
|
|
group.throughput(Throughput::Elements(1));
|
|
|
|
let config = optimization::parallelism::TensorParallelConfig {
|
|
world_size: 4,
|
|
rank: 0,
|
|
device_ids: vec![0, 1, 2, 3],
|
|
};
|
|
|
|
let backend = Arc::new(optimization::parallelism::NCCLBackend::new());
|
|
let mut parallel_system = optimization::parallelism::ModelParallelism::new(config, backend);
|
|
|
|
let model = Arc::new(MockModelInterface::new("/tmp/mock").unwrap());
|
|
parallel_system.add_local_model(model);
|
|
|
|
group.bench_function("tensor_parallel_forward", |b| {
|
|
b.to_async(&rt).iter(|| async {
|
|
let input_ids = rtx_tensor::Tensor::from_slice(&[1u32, 2, 3, 4, 5], &[1, 5]).unwrap();
|
|
let result = parallel_system
|
|
.parallel_forward(&input_ids, None)
|
|
.await
|
|
.unwrap();
|
|
black_box(result);
|
|
});
|
|
});
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_batch_optimization(c: &mut Criterion) {
|
|
let rt = tokio::runtime::Runtime::new().unwrap();
|
|
|
|
let mut group = c.benchmark_group("batch_optimization");
|
|
|
|
for (strategy_name, strategy) in [
|
|
("FIFO", serving::batch::SchedulingStrategy::FIFO),
|
|
("SJF", serving::batch::SchedulingStrategy::ShortestJobFirst),
|
|
("Priority", serving::batch::SchedulingStrategy::Priority),
|
|
("Adaptive", serving::batch::SchedulingStrategy::Adaptive),
|
|
]
|
|
.iter()
|
|
{
|
|
group.bench_with_input(
|
|
BenchmarkId::new("strategy", strategy_name),
|
|
strategy,
|
|
|b, strategy| {
|
|
b.to_async(&rt).iter(|| async {
|
|
let model = Arc::new(MockModelInterface::new("/tmp/mock").unwrap());
|
|
let gen_config = GenerationConfig::default().max_length(50);
|
|
|
|
let batch_config = serving::batch::BatchConfig {
|
|
max_batch_size: 8,
|
|
scheduling_strategy: strategy.clone(),
|
|
dynamic_batching: true,
|
|
..Default::default()
|
|
};
|
|
|
|
let generator = serving::batch::BatchGenerator::with_batch_config(
|
|
model,
|
|
gen_config,
|
|
batch_config,
|
|
)
|
|
.unwrap();
|
|
|
|
let prompts = (0..8).map(|i| format!("Test prompt {}", i)).collect();
|
|
|
|
let results = generator.generate_batch(prompts).await.unwrap();
|
|
black_box(results);
|
|
});
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_quality_assessment(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("quality_assessment");
|
|
group.throughput(Throughput::Elements(1));
|
|
|
|
let test_texts = [
|
|
"Short text.",
|
|
"Medium length text with several words and some variety in vocabulary.",
|
|
"Long text with multiple sentences containing diverse vocabulary, complex structures, and varied linguistic patterns that should trigger comprehensive quality assessment.",
|
|
];
|
|
|
|
for (i, text) in test_texts.iter().enumerate() {
|
|
group.bench_with_input(
|
|
BenchmarkId::new("words", text.split_whitespace().count()),
|
|
text,
|
|
|b, text| {
|
|
b.iter(|| {
|
|
// Mock quality assessment
|
|
let quality_score = calculate_mock_quality(black_box(text));
|
|
black_box(quality_score);
|
|
});
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn calculate_mock_quality(text: &str) -> f32 {
|
|
let words = text.split_whitespace().collect::<Vec<_>>();
|
|
let unique_words: std::collections::HashSet<_> = words.iter().collect();
|
|
|
|
let diversity = unique_words.len() as f32 / words.len().max(1) as f32;
|
|
let length_factor = (words.len() as f32 / 50.0).min(1.0);
|
|
|
|
(diversity * 0.7 + length_factor * 0.3).clamp(0.0, 1.0)
|
|
}
|
|
|
|
fn bench_end_to_end_pipeline(c: &mut Criterion) {
|
|
let rt = tokio::runtime::Runtime::new().unwrap();
|
|
|
|
let mut group = c.benchmark_group("end_to_end_pipeline");
|
|
group.throughput(Throughput::Elements(1));
|
|
|
|
group.bench_function("complete_nlg_pipeline", |b| {
|
|
b.to_async(&rt).iter(|| async {
|
|
// 1. Template rendering
|
|
let mut template_manager = serving::templates::TemplateManager::new();
|
|
let mut values = HashMap::new();
|
|
values.insert(
|
|
"text".to_string(),
|
|
"AI is transforming the world".to_string(),
|
|
);
|
|
values.insert("max_sentences".to_string(), "3".to_string());
|
|
|
|
let prompt = template_manager
|
|
.render_template("summarize", &values)
|
|
.unwrap();
|
|
|
|
// 2. Text generation with constraints
|
|
let model = Arc::new(MockModelInterface::new("/tmp/mock").unwrap());
|
|
let config = GenerationConfig::nucleus_sampling()
|
|
.max_length(100)
|
|
.temperature(0.8);
|
|
|
|
let mut generator = TextGenerator::with_model(model.clone(), config).unwrap();
|
|
let generated = generator.generate(&prompt).unwrap();
|
|
|
|
// 3. Quality assessment
|
|
let quality = calculate_mock_quality(&generated.text);
|
|
|
|
// 4. Stream output
|
|
let streaming_config = GenerationConfig::default().max_length(50);
|
|
let streaming_generator =
|
|
serving::StreamingGenerator::with_model(model, streaming_config).unwrap();
|
|
|
|
let mut stream = streaming_generator.generate_stream(&prompt).await.unwrap();
|
|
let mut tokens = Vec::new();
|
|
|
|
// Collect a few tokens
|
|
for _ in 0..5 {
|
|
if let Some(token) = stream.next().await {
|
|
tokens.push(token.unwrap());
|
|
} else {
|
|
break;
|
|
}
|
|
}
|
|
|
|
black_box((generated, quality, tokens));
|
|
});
|
|
});
|
|
|
|
group.finish();
|
|
}
|
|
|
|
criterion_group!(
|
|
integration_benches,
|
|
bench_streaming_generation,
|
|
bench_constrained_generation,
|
|
bench_controllable_generation,
|
|
bench_template_rendering,
|
|
bench_parallel_generation,
|
|
bench_batch_optimization,
|
|
bench_quality_assessment,
|
|
bench_end_to_end_pipeline
|
|
);
|
|
|
|
criterion_main!(integration_benches);
|