macrocc
69b37ced07
- 将项目拆分为三个 crate：minicpm-core（核心模型）、minicpm-convert（转换功能）、minicpm-inference（推理功能） - 添加两个示例：minimal-inference（最小推理）和 convert（模型转换） - 转换后自动拷贝 config.json 和 tokenizer.json 到 model 目录 - 更新 README 说明 workspace 结构和使用方式
83 lines
2.6 KiB
Rust
83 lines
2.6 KiB
Rust
use burn::tensor::backend::Backend;
|
|
use burn::tensor::{Float, Tensor};
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct RoPE {
|
|
pub dim: usize,
|
|
pub theta: f64,
|
|
pub max_seq_len: usize,
|
|
cos_cache: Vec<f32>,
|
|
sin_cache: Vec<f32>,
|
|
half_dim: usize,
|
|
}
|
|
|
|
impl RoPE {
|
|
pub fn new(dim: usize, theta: f64, max_seq_len: usize) -> Self {
|
|
let half_dim = dim / 2;
|
|
|
|
let mut inv_freqs = Vec::with_capacity(half_dim);
|
|
for i in 0..half_dim {
|
|
let freq = 1.0 / (theta.powf((2.0 * i as f64) / dim as f64));
|
|
inv_freqs.push(freq as f32);
|
|
}
|
|
|
|
let mut cos_cache = Vec::with_capacity(max_seq_len * half_dim);
|
|
let mut sin_cache = Vec::with_capacity(max_seq_len * half_dim);
|
|
|
|
for pos in 0..max_seq_len {
|
|
for &freq in &inv_freqs {
|
|
let angle = pos as f32 * freq;
|
|
cos_cache.push(angle.cos());
|
|
sin_cache.push(angle.sin());
|
|
}
|
|
}
|
|
|
|
Self {
|
|
dim,
|
|
theta,
|
|
max_seq_len,
|
|
cos_cache,
|
|
sin_cache,
|
|
half_dim,
|
|
}
|
|
}
|
|
|
|
pub fn get_cache(&self, offset: usize, len: usize) -> (Vec<f32>, Vec<f32>) {
|
|
let cos = self.cos_cache[offset * self.half_dim..(offset + len) * self.half_dim].to_vec();
|
|
let sin = self.sin_cache[offset * self.half_dim..(offset + len) * self.half_dim].to_vec();
|
|
(cos, sin)
|
|
}
|
|
|
|
pub fn apply<B: Backend>(
|
|
&self,
|
|
x: Tensor<B, 4, Float>,
|
|
offset: usize,
|
|
) -> Tensor<B, 4, Float> {
|
|
let device = x.device();
|
|
let shape = x.shape();
|
|
let dims: [usize; 4] = shape.dims();
|
|
let batch = dims[0];
|
|
let heads = dims[1];
|
|
let seq_len = dims[2];
|
|
let dim = dims[3];
|
|
let half_dim = self.half_dim;
|
|
|
|
let cos_vals = &self.cos_cache[offset * half_dim..(offset + seq_len) * half_dim];
|
|
let sin_vals = &self.sin_cache[offset * half_dim..(offset + seq_len) * half_dim];
|
|
|
|
let cos = Tensor::<B, 1, Float>::from_floats(cos_vals, &device)
|
|
.reshape([1, 1, seq_len, half_dim])
|
|
.expand([batch, heads, seq_len, half_dim]);
|
|
let sin = Tensor::<B, 1, Float>::from_floats(sin_vals, &device)
|
|
.reshape([1, 1, seq_len, half_dim])
|
|
.expand([batch, heads, seq_len, half_dim]);
|
|
|
|
let x1 = x.clone().slice([0..batch, 0..heads, 0..seq_len, 0..half_dim]);
|
|
let x2 = x.slice([0..batch, 0..heads, 0..seq_len, half_dim..dim]);
|
|
|
|
let x1_rot = x1.clone() * cos.clone() - x2.clone() * sin.clone();
|
|
let x2_rot = x2 * cos + x1 * sin;
|
|
|
|
Tensor::cat(vec![x1_rot, x2_rot], 3)
|
|
}
|
|
} |