Struct LinearProbe 

pub struct LinearProbe {
    pub weight: Tensor,
    pub bias: Tensor,
    /* private fields */
}

Linear probe: frozen embeddings + trainable linear head (CLF-002).

Trains on pre-extracted embeddings (not raw token IDs), making training complete in seconds rather than minutes.

Fields

weight: Tensor

Linear weight [hidden_size, num_classes] flattened row-major

bias: Tensor

Bias [num_classes]

Implementations

impl LinearProbe

pub fn new(hidden_size: usize, num_classes: usize) -> Self

Create with Xavier initialization.

pub fn forward(&self, embedding: &Tensor) -> Tensor

Forward pass: embedding → logits.

Arguments

• embedding - Pre-extracted [CLS] embedding [hidden_size]

Returns

Logits tensor [num_classes]

pub fn predict_probs(&self, embedding: &Tensor) -> Vec<f32>

Predict class probabilities via softmax.

pub fn predict(&self, embedding: &Tensor) -> usize

Predict class index (argmax of logits).

pub fn train( &mut self, embeddings: &[Vec<f32>], labels: &[usize], epochs: usize, learning_rate: f32, class_weights: Option<&[f32]>, ) -> f32

Train on pre-extracted embeddings using SGD.

Arguments

• embeddings - List of pre-extracted [CLS] embeddings (each len=hidden_size)
• labels - Corresponding class labels
• epochs - Number of training epochs
• learning_rate - SGD learning rate
• class_weights - Optional per-class loss weights for imbalance

Returns

Final training loss

pub fn num_parameters(&self) -> usize

Get total trainable parameter count (CLF-002: 1,538 for binary CodeBERT).

pub fn num_classes(&self) -> usize

Get number of classes.