use super::neural_types::*;
use super::neural_utilities::*;
use scirs2_core::ndarray::s;
pub fn sinusoidal_positional_encoding(seq_len: usize, d_model: usize) -> Array2<f64> {
let mut encoding = Array2::zeros((seq_len, d_model));
for pos in 0..seq_len {
for i in 0..d_model {
let angle = pos as f64 / (10000.0_f64).powf((2 * (i / 2)) as f64 / d_model as f64);
if i % 2 == 0 {
encoding[(pos, i)] = angle.sin();
} else {
encoding[(pos, i)] = angle.cos();
}
}
}
encoding
}
pub fn learned_positional_encoding(seq_len: usize, d_model: usize, random_state: Option<u64>) -> Array2<f64> {
xavier_normal_init(seq_len, d_model, random_state)
}
pub fn relative_positional_encoding(seq_len: usize, d_model: usize, max_relative_position: usize) -> Array3<f64> {
let mut encoding = Array3::zeros((seq_len, seq_len, d_model));
for i in 0..seq_len {
for j in 0..seq_len {
let relative_pos = ((i as i32 - j as i32).abs() as usize).min(max_relative_position);
for k in 0..d_model {
let angle = relative_pos as f64 / (10000.0_f64).powf((2 * (k / 2)) as f64 / d_model as f64);
if k % 2 == 0 {
encoding[(i, j, k)] = angle.sin();
} else {
encoding[(i, j, k)] = angle.cos();
}
}
}
}
encoding
}
pub fn rotary_positional_encoding(seq_len: usize, d_model: usize) -> Array2<f64> {
assert_eq!(d_model % 2, 0, "d_model must be even for rotary encoding");
let mut encoding = Array2::zeros((seq_len, d_model));
for pos in 0..seq_len {
for i in (0..d_model).step_by(2) {
let theta = pos as f64 / (10000.0_f64).powf(i as f64 / d_model as f64);
encoding[(pos, i)] = theta.cos();
encoding[(pos, i + 1)] = theta.sin();
}
}
encoding
}
pub fn apply_rotary_encoding(x: &Array2<f64>, pos_encoding: &Array2<f64>) -> Array2<f64> {
let (seq_len, d_model) = x.dim();
let mut result = Array2::zeros((seq_len, d_model));
for i in 0..seq_len {
for j in (0..d_model).step_by(2) {
let cos_theta = pos_encoding[(i, j)];
let sin_theta = pos_encoding[(i, j + 1)];
let x1 = x[(i, j)];
let x2 = x[(i, j + 1)];
result[(i, j)] = x1 * cos_theta - x2 * sin_theta;
result[(i, j + 1)] = x1 * sin_theta + x2 * cos_theta;
}
}
result
}
pub fn alibi_positional_bias(seq_len: usize, n_heads: usize) -> Array3<f64> {
let mut bias = Array3::zeros((n_heads, seq_len, seq_len));
let slopes = alibi_get_slopes(n_heads);
for h in 0..n_heads {
for i in 0..seq_len {
for j in 0..seq_len {
if j <= i {
bias[(h, i, j)] = -slopes[h] * (i - j) as f64;
} else {
bias[(h, i, j)] = f64::NEG_INFINITY;
}
}
}
}
bias
}
fn alibi_get_slopes(n_heads: usize) -> Vec<f64> {
let mut slopes = Vec::new();
if n_heads.is_power_of_two() {
let start = 2.0_f64.powf(-8.0 / n_heads as f64);
let ratio = 2.0_f64.powf(-1.0);
for i in 0..n_heads {
slopes.push(start * ratio.powf(i as f64));
}
} else {
let closest_power_of_2 = 2_usize.pow((n_heads as f64).log2().floor() as u32);
let extra_heads = n_heads - closest_power_of_2;
let base_slopes = alibi_get_slopes(closest_power_of_2);
slopes.extend_from_slice(&base_slopes);
let start = 2.0_f64.powf(-8.0 / (2 * closest_power_of_2) as f64);
let ratio = 2.0_f64.powf(-2.0);
for i in 0..extra_heads {
slopes.push(start * ratio.powf(i as f64));
}
}
slopes
}
pub struct PositionalEmbedding {
pub encoding_type: PositionalEncodingType,
pub seq_len: usize,
pub d_model: usize,
pub learned_embeddings: Option<Array2<f64>>,
pub max_relative_position: usize,
pub dropout_rate: f64,
}
#[derive(Debug, Clone, Copy)]
pub enum PositionalEncodingType {
Sinusoidal,
Learned,
Relative,
Rotary,
ALiBi,
None,
}
impl PositionalEmbedding {
pub fn new(
encoding_type: PositionalEncodingType,
seq_len: usize,
d_model: usize,
random_state: Option<u64>,
) -> Self {
let learned_embeddings = match encoding_type {
PositionalEncodingType::Learned => Some(learned_positional_encoding(seq_len, d_model, random_state)),
_ => None,
};
Self {
encoding_type,
seq_len,
d_model,
learned_embeddings,
max_relative_position: 32,
dropout_rate: 0.1,
}
}
pub fn with_max_relative_position(mut self, max_relative_position: usize) -> Self {
self.max_relative_position = max_relative_position;
self
}
pub fn with_dropout(mut self, dropout_rate: f64) -> Self {
self.dropout_rate = dropout_rate;
self
}
pub fn get_encoding(&self, actual_seq_len: usize) -> Array2<f64> {
let effective_seq_len = actual_seq_len.min(self.seq_len);
match self.encoding_type {
PositionalEncodingType::Sinusoidal => {
sinusoidal_positional_encoding(effective_seq_len, self.d_model)
},
PositionalEncodingType::Learned => {
if let Some(ref embeddings) = self.learned_embeddings {
embeddings.slice(s![..effective_seq_len, ..]).to_owned()
} else {
Array2::zeros((effective_seq_len, self.d_model))
}
},
PositionalEncodingType::Rotary => {
rotary_positional_encoding(effective_seq_len, self.d_model)
},
_ => Array2::zeros((effective_seq_len, self.d_model)),
}
}
pub fn apply_encoding(&self, x: &Array2<f64>) -> Array2<f64> {
let (seq_len, d_model) = x.dim();
if d_model != self.d_model {
panic!("Input d_model {} does not match expected d_model {}", d_model, self.d_model);
}
match self.encoding_type {
PositionalEncodingType::Sinusoidal | PositionalEncodingType::Learned => {
let encoding = self.get_encoding(seq_len);
x + &encoding
},
PositionalEncodingType::Rotary => {
let encoding = self.get_encoding(seq_len);
apply_rotary_encoding(x, &encoding)
},
PositionalEncodingType::None => x.clone(),
_ => x.clone(),
}
}
pub fn get_relative_encoding(&self, seq_len: usize) -> Array3<f64> {
relative_positional_encoding(seq_len, self.d_model, self.max_relative_position)
}
pub fn get_alibi_bias(&self, seq_len: usize, n_heads: usize) -> Array3<f64> {
alibi_positional_bias(seq_len, n_heads)
}
pub fn apply_dropout(&self, x: &Array2<f64>) -> Array2<f64> {
if self.dropout_rate == 0.0 {
return x.clone();
}
let mut rng = Random::seed_from_u64(42);
let mut result = x.clone();
for elem in result.iter_mut() {
if rng.random_range(0.0..1.0) < self.dropout_rate {
*elem = 0.0;
} else {
*elem /= 1.0 - self.dropout_rate;
}
}
result
}
}
pub fn create_position_ids(seq_len: usize, batch_size: usize) -> Array2<usize> {
let mut position_ids = Array2::zeros((batch_size, seq_len));
for i in 0..batch_size {
for j in 0..seq_len {
position_ids[(i, j)] = j;
}
}
position_ids
}
pub fn create_token_type_ids(seq_lengths: &[usize], max_len: usize) -> Array2<usize> {
let batch_size = seq_lengths.len();
let mut token_type_ids = Array2::zeros((batch_size, max_len));
for (i, &length) in seq_lengths.iter().enumerate() {
for j in 0..length.min(max_len) {
token_type_ids[(i, j)] = if j < length / 2 { 0 } else { 1 };
}
}
token_type_ids
}
pub fn interpolate_positional_encoding(
original_encoding: &Array2<f64>,
new_seq_len: usize,
) -> Array2<f64> {
let (original_seq_len, d_model) = original_encoding.dim();
if new_seq_len == original_seq_len {
return original_encoding.clone();
}
let mut interpolated = Array2::zeros((new_seq_len, d_model));
for i in 0..new_seq_len {
let position = (i as f64 * (original_seq_len - 1) as f64) / (new_seq_len - 1) as f64;
let lower_idx = position.floor() as usize;
let upper_idx = (lower_idx + 1).min(original_seq_len - 1);
let weight = position - position.floor();
for j in 0..d_model {
if lower_idx == upper_idx {
interpolated[(i, j)] = original_encoding[(lower_idx, j)];
} else {
interpolated[(i, j)] = (1.0 - weight) * original_encoding[(lower_idx, j)]
+ weight * original_encoding[(upper_idx, j)];
}
}
}
interpolated
}
pub fn extrapolate_positional_encoding(
original_encoding: &Array2<f64>,
new_seq_len: usize,
) -> Array2<f64> {
let (original_seq_len, d_model) = original_encoding.dim();
if new_seq_len <= original_seq_len {
return interpolate_positional_encoding(original_encoding, new_seq_len);
}
let mut extrapolated = Array2::zeros((new_seq_len, d_model));
for j in 0..d_model {
let last_value = original_encoding[(original_seq_len - 1, j)];
let second_last_value = if original_seq_len > 1 {
original_encoding[(original_seq_len - 2, j)]
} else {
0.0
};
let diff = last_value - second_last_value;
for i in 0..original_seq_len {
extrapolated[(i, j)] = original_encoding[(i, j)];
}
for i in original_seq_len..new_seq_len {
let steps = (i - original_seq_len + 1) as f64;
extrapolated[(i, j)] = last_value + diff * steps;
}
}
extrapolated
}
pub fn compute_positional_similarity(
pos1: usize,
pos2: usize,
encoding: &Array2<f64>,
) -> f64 {
if pos1 >= encoding.nrows() || pos2 >= encoding.nrows() {
return 0.0;
}
let vec1 = encoding.row(pos1);
let vec2 = encoding.row(pos2);
cosine_similarity(&vec1.to_owned(), &vec2.to_owned())
}
pub fn positional_encoding_temperature_scaling(
encoding: &Array2<f64>,
temperature: f64,
) -> Array2<f64> {
encoding / temperature
}
pub fn adaptive_positional_encoding(
seq_len: usize,
d_model: usize,
adaptation_factor: f64,
) -> Array2<f64> {
let base_encoding = sinusoidal_positional_encoding(seq_len, d_model);
let mut adaptive_encoding = Array2::zeros((seq_len, d_model));
for i in 0..seq_len {
let position_weight = 1.0 + adaptation_factor * (i as f64 / seq_len as f64);
for j in 0..d_model {
adaptive_encoding[(i, j)] = base_encoding[(i, j)] * position_weight;
}
}
adaptive_encoding
}
pub fn hierarchical_positional_encoding(
seq_len: usize,
d_model: usize,
hierarchy_levels: usize,
) -> Array2<f64> {
let mut encoding = Array2::zeros((seq_len, d_model));
let dims_per_level = d_model / hierarchy_levels;
for level in 0..hierarchy_levels {
let scale = 2_usize.pow(level as u32);
let start_dim = level * dims_per_level;
let end_dim = ((level + 1) * dims_per_level).min(d_model);
for pos in 0..seq_len {
for i in start_dim..end_dim {
let dim_in_level = i - start_dim;
let angle = (pos / scale) as f64 / (10000.0_f64).powf((2 * (dim_in_level / 2)) as f64 / dims_per_level as f64);
if dim_in_level % 2 == 0 {
encoding[(pos, i)] = angle.sin();
} else {
encoding[(pos, i)] = angle.cos();
}
}
}
}
encoding
}
pub fn contextual_positional_encoding(
context_vectors: &Array2<f64>,
d_model: usize,
random_state: Option<u64>,
) -> Array2<f64> {
let seq_len = context_vectors.nrows();
let context_dim = context_vectors.ncols();
let projection_matrix = xavier_normal_init(context_dim, d_model, random_state);
let base_encoding = sinusoidal_positional_encoding(seq_len, d_model);
let projected_context = context_vectors.dot(&projection_matrix);
&base_encoding + &projected_context
}