pub fn generate_span_indices(text_len: usize, max_width: usize) -> Vec<[usize; 2]> {
let mut indices = Vec::with_capacity(text_len * max_width);
for i in 0..text_len {
for w in 0..max_width {
let end = (i + w).min(text_len.saturating_sub(1));
indices.push([i, end]);
}
}
indices
}
pub fn generate_batched_span_indices(
text_lengths: &[usize],
max_text_len: usize,
max_width: usize,
) -> Vec<[usize; 2]> {
let n_spans = max_text_len * max_width;
let batch_size = text_lengths.len();
let mut result = vec![[0usize; 2]; batch_size * n_spans];
for (b, &tl) in text_lengths.iter().enumerate() {
for i in 0..max_text_len {
for w in 0..max_width {
let idx = i * max_width + w;
let end = i + w;
if end < tl {
result[b * n_spans + idx] = [i, end];
}
}
}
}
result
}