mirror of
https://github.com/ruvnet/RuView
synced 2026-08-07 20:01:43 +00:00
Merge commit 'd803bfe2b1fe7f5e219e50ac20d6801a0a58ac75' as 'vendor/ruvector'
This commit is contained in:
ruv
@@ -0,0 +1,333 @@
|
||||
//! Embedding operations for ESP32
|
||||
//!
|
||||
//! Provides efficient token embedding lookup and positional encoding.
|
||||
|
||||
use heapless::Vec as HVec;
|
||||
|
||||
/// Maximum embedding dimension
|
||||
pub const MAX_EMBED_DIM: usize = 128;
|
||||
/// Maximum vocabulary size for stack allocation
|
||||
pub const MAX_VOCAB: usize = 2048;
|
||||
|
||||
/// Embedding table with INT8 quantization
|
||||
pub struct EmbeddingTable<const VOCAB: usize, const DIM: usize> {
|
||||
/// Flattened embedding weights [VOCAB * DIM]
|
||||
weights: HVec<i8, { 64 * 1024 }>, // Max 64KB
|
||||
/// Vocabulary size
|
||||
vocab_size: usize,
|
||||
/// Embedding dimension
|
||||
embed_dim: usize,
|
||||
/// Scale factor for dequantization
|
||||
scale: f32,
|
||||
}
|
||||
|
||||
impl<const VOCAB: usize, const DIM: usize> EmbeddingTable<VOCAB, DIM> {
|
||||
/// Create new embedding table from weights
|
||||
pub fn new(weights: &[i8], vocab_size: usize, embed_dim: usize) -> crate::Result<Self> {
|
||||
if weights.len() != vocab_size * embed_dim {
|
||||
return Err(crate::Error::InvalidModel("Weight size mismatch"));
|
||||
}
|
||||
|
||||
let mut table_weights = HVec::new();
|
||||
for &w in weights {
|
||||
table_weights.push(w).map_err(|_| crate::Error::BufferOverflow)?;
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
weights: table_weights,
|
||||
vocab_size,
|
||||
embed_dim,
|
||||
scale: 1.0 / 127.0,
|
||||
})
|
||||
}
|
||||
|
||||
/// Create random embedding table for testing
|
||||
pub fn random(vocab_size: usize, embed_dim: usize, seed: u32) -> crate::Result<Self> {
|
||||
let mut weights = HVec::new();
|
||||
let mut rng_state = seed;
|
||||
|
||||
for _ in 0..(vocab_size * embed_dim) {
|
||||
rng_state = rng_state.wrapping_mul(1103515245).wrapping_add(12345);
|
||||
let val = ((rng_state >> 16) & 0xFF) as i8;
|
||||
weights.push(val).map_err(|_| crate::Error::BufferOverflow)?;
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
weights,
|
||||
vocab_size,
|
||||
embed_dim,
|
||||
scale: 1.0 / 127.0,
|
||||
})
|
||||
}
|
||||
|
||||
/// Look up embedding for a token
|
||||
#[inline]
|
||||
pub fn lookup(&self, token_id: u16, output: &mut [i8]) -> crate::Result<()> {
|
||||
let id = token_id as usize;
|
||||
if id >= self.vocab_size {
|
||||
return Err(crate::Error::InvalidModel("Token ID out of range"));
|
||||
}
|
||||
|
||||
let start = id * self.embed_dim;
|
||||
let end = start + self.embed_dim;
|
||||
|
||||
if output.len() < self.embed_dim {
|
||||
return Err(crate::Error::BufferOverflow);
|
||||
}
|
||||
|
||||
output[..self.embed_dim].copy_from_slice(&self.weights[start..end]);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Look up embedding and add to existing buffer (for accumulation)
|
||||
#[inline]
|
||||
pub fn lookup_add(&self, token_id: u16, output: &mut [i32]) -> crate::Result<()> {
|
||||
let id = token_id as usize;
|
||||
if id >= self.vocab_size {
|
||||
return Err(crate::Error::InvalidModel("Token ID out of range"));
|
||||
}
|
||||
|
||||
let start = id * self.embed_dim;
|
||||
|
||||
for i in 0..self.embed_dim {
|
||||
output[i] += self.weights[start + i] as i32;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Memory size in bytes
|
||||
pub fn memory_size(&self) -> usize {
|
||||
self.weights.len()
|
||||
}
|
||||
}
|
||||
|
||||
/// Rotary Position Embedding (RoPE) for ESP32
|
||||
///
|
||||
/// Uses fixed-point arithmetic for sin/cos computation.
|
||||
pub struct RotaryEmbedding {
|
||||
/// Dimension (must be even)
|
||||
dim: usize,
|
||||
/// Base frequency
|
||||
base: u32,
|
||||
/// Precomputed sin values (fixed-point, scaled by 128)
|
||||
sin_cache: [i8; MAX_EMBED_DIM],
|
||||
/// Precomputed cos values (fixed-point, scaled by 128)
|
||||
cos_cache: [i8; MAX_EMBED_DIM],
|
||||
/// Maximum cached position
|
||||
max_cached_pos: usize,
|
||||
}
|
||||
|
||||
impl RotaryEmbedding {
|
||||
/// Create new RoPE with given dimension
|
||||
pub fn new(dim: usize, base: u32) -> Self {
|
||||
Self {
|
||||
dim,
|
||||
base,
|
||||
sin_cache: [0i8; MAX_EMBED_DIM],
|
||||
cos_cache: [0i8; MAX_EMBED_DIM],
|
||||
max_cached_pos: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Update cache for new position
|
||||
pub fn update_cache(&mut self, pos: usize) {
|
||||
if pos <= self.max_cached_pos {
|
||||
return;
|
||||
}
|
||||
|
||||
// Compute frequency for each dimension pair
|
||||
for i in 0..(self.dim / 2) {
|
||||
// freq = 1 / (base^(2i/dim))
|
||||
// For INT8, we approximate using lookup table or simple formula
|
||||
|
||||
// Simplified: use position-dependent rotation
|
||||
// angle = pos / (base^(i / (dim/2)))
|
||||
let freq_scale = ((i * 256) / (self.dim / 2)) as u32;
|
||||
let angle = ((pos as u32 * 256) / (self.base + freq_scale)) as i32;
|
||||
|
||||
// Approximate sin/cos using polynomial
|
||||
// sin(x) ≈ x - x³/6 for small x (scaled)
|
||||
// cos(x) ≈ 1 - x²/2 for small x (scaled)
|
||||
let x = (angle % 256) as i32 - 128; // Center around 0
|
||||
|
||||
// Simple quadrant-based approximation
|
||||
let sin_val = (x * 127 / 128).clamp(-127, 127) as i8;
|
||||
let cos_val = ((128 - x.abs()) * 127 / 128).clamp(-127, 127) as i8;
|
||||
|
||||
self.sin_cache[i] = sin_val;
|
||||
self.cos_cache[i] = cos_val;
|
||||
self.sin_cache[i + self.dim / 2] = sin_val;
|
||||
self.cos_cache[i + self.dim / 2] = cos_val;
|
||||
}
|
||||
|
||||
self.max_cached_pos = pos;
|
||||
}
|
||||
|
||||
/// Apply rotary embedding to query/key vectors
|
||||
#[inline]
|
||||
pub fn apply(&self, x: &mut [i8], _pos: usize) {
|
||||
let half_dim = self.dim / 2;
|
||||
|
||||
// Process pairs of dimensions
|
||||
for i in 0..half_dim {
|
||||
let x1 = x[i] as i32;
|
||||
let x2 = x[i + half_dim] as i32;
|
||||
|
||||
let sin = self.sin_cache[i] as i32;
|
||||
let cos = self.cos_cache[i] as i32;
|
||||
|
||||
// Rotation: [cos, -sin; sin, cos] @ [x1, x2]
|
||||
let new_x1 = (x1 * cos - x2 * sin) >> 7;
|
||||
let new_x2 = (x1 * sin + x2 * cos) >> 7;
|
||||
|
||||
x[i] = new_x1.clamp(-128, 127) as i8;
|
||||
x[i + half_dim] = new_x2.clamp(-128, 127) as i8;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Simple positional encoding using learned embeddings
|
||||
pub struct LearnedPositionalEmbedding<const MAX_LEN: usize, const DIM: usize> {
|
||||
/// Position embeddings [MAX_LEN * DIM]
|
||||
embeddings: HVec<i8, { 8 * 1024 }>, // Max 8KB for positions
|
||||
/// Maximum sequence length
|
||||
max_len: usize,
|
||||
/// Embedding dimension
|
||||
dim: usize,
|
||||
}
|
||||
|
||||
impl<const MAX_LEN: usize, const DIM: usize> LearnedPositionalEmbedding<MAX_LEN, DIM> {
|
||||
/// Create random positional embeddings
|
||||
pub fn random(max_len: usize, dim: usize, seed: u32) -> crate::Result<Self> {
|
||||
let mut embeddings = HVec::new();
|
||||
let mut rng_state = seed;
|
||||
|
||||
for _ in 0..(max_len * dim) {
|
||||
rng_state = rng_state.wrapping_mul(1103515245).wrapping_add(12345);
|
||||
// Smaller values for positional embeddings
|
||||
let val = (((rng_state >> 16) & 0x3F) as i8) - 32;
|
||||
embeddings.push(val).map_err(|_| crate::Error::BufferOverflow)?;
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
embeddings,
|
||||
max_len,
|
||||
dim,
|
||||
})
|
||||
}
|
||||
|
||||
/// Add positional embedding to input
|
||||
#[inline]
|
||||
pub fn add_to(&self, input: &mut [i8], pos: usize) -> crate::Result<()> {
|
||||
if pos >= self.max_len {
|
||||
return Err(crate::Error::BufferOverflow);
|
||||
}
|
||||
|
||||
let start = pos * self.dim;
|
||||
for i in 0..self.dim {
|
||||
let sum = input[i] as i32 + self.embeddings[start + i] as i32;
|
||||
input[i] = sum.clamp(-128, 127) as i8;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Memory size in bytes
|
||||
pub fn memory_size(&self) -> usize {
|
||||
self.embeddings.len()
|
||||
}
|
||||
}
|
||||
|
||||
/// Byte-Pair Encoding tokenizer (simplified)
|
||||
///
|
||||
/// For ESP32, we use a simple character-level or small vocabulary tokenizer.
|
||||
pub struct SimpleTokenizer {
|
||||
/// Character to token ID mapping
|
||||
char_to_id: [u16; 256],
|
||||
/// Token ID to character mapping
|
||||
id_to_char: [u8; 256],
|
||||
/// Vocabulary size
|
||||
vocab_size: usize,
|
||||
}
|
||||
|
||||
impl SimpleTokenizer {
|
||||
/// Create ASCII tokenizer (vocabulary = 128)
|
||||
pub fn ascii() -> Self {
|
||||
let mut char_to_id = [0u16; 256];
|
||||
let mut id_to_char = [0u8; 256];
|
||||
|
||||
for i in 0..128 {
|
||||
char_to_id[i] = i as u16;
|
||||
id_to_char[i] = i as u8;
|
||||
}
|
||||
|
||||
// Map non-ASCII to UNK (127)
|
||||
for i in 128..256 {
|
||||
char_to_id[i] = 127;
|
||||
}
|
||||
|
||||
Self {
|
||||
char_to_id,
|
||||
id_to_char,
|
||||
vocab_size: 128,
|
||||
}
|
||||
}
|
||||
|
||||
/// Tokenize a string
|
||||
pub fn encode(&self, text: &str) -> HVec<u16, 128> {
|
||||
let mut tokens = HVec::new();
|
||||
for byte in text.bytes() {
|
||||
let _ = tokens.push(self.char_to_id[byte as usize]);
|
||||
}
|
||||
tokens
|
||||
}
|
||||
|
||||
/// Decode tokens to string
|
||||
pub fn decode(&self, tokens: &[u16]) -> HVec<u8, 128> {
|
||||
let mut chars = HVec::new();
|
||||
for &token in tokens {
|
||||
if (token as usize) < self.vocab_size {
|
||||
let _ = chars.push(self.id_to_char[token as usize]);
|
||||
}
|
||||
}
|
||||
chars
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_embedding_lookup() {
|
||||
let embed: EmbeddingTable<256, 64> = EmbeddingTable::random(256, 64, 42).unwrap();
|
||||
|
||||
let mut output = [0i8; 64];
|
||||
embed.lookup(10, &mut output).unwrap();
|
||||
|
||||
// Should be non-zero
|
||||
assert!(output.iter().any(|&x| x != 0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rotary_embedding() {
|
||||
let mut rope = RotaryEmbedding::new(32, 10000);
|
||||
rope.update_cache(10);
|
||||
|
||||
let mut x = [64i8; 32];
|
||||
rope.apply(&mut x, 5);
|
||||
|
||||
// Values should change after rotation
|
||||
assert!(x.iter().any(|&v| v != 64));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tokenizer() {
|
||||
let tokenizer = SimpleTokenizer::ascii();
|
||||
|
||||
let tokens = tokenizer.encode("Hello");
|
||||
assert_eq!(tokens.len(), 5);
|
||||
|
||||
let decoded = tokenizer.decode(&tokens);
|
||||
assert_eq!(&decoded[..], b"Hello");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user