ghostflow_nn/

clip.rs

//! CLIP (Contrastive Language-Image Pre-training)
//!
//! Implements CLIP architecture:
//! - Vision Transformer or ResNet for image encoding
//! - Transformer for text encoding
//! - Contrastive learning objective
//! - Zero-shot classification
//! - Text-image similarity

use ghostflow_core::Tensor;
use crate::linear::Linear;
use crate::vision_transformer::{VisionTransformer, ViTConfig};
use crate::Module;

/// CLIP configuration
#[derive(Debug, Clone)]
pub struct CLIPConfig {
    /// Embedding dimension (shared between vision and text)
    pub embed_dim: usize,
    /// Vision config
    pub vision_config: CLIPVisionConfig,
    /// Text config
    pub text_config: CLIPTextConfig,
    /// Logit scale initialization
    pub logit_scale_init_value: f32,
}

/// CLIP Vision configuration
#[derive(Debug, Clone)]
pub struct CLIPVisionConfig {
    /// Image size
    pub image_size: usize,
    /// Patch size
    pub patch_size: usize,
    /// Hidden size
    pub hidden_size: usize,
    /// Number of layers
    pub num_layers: usize,
    /// Number of attention heads
    pub num_heads: usize,
    /// MLP ratio
    pub mlp_ratio: usize,
}

/// CLIP Text configuration
#[derive(Debug, Clone)]
pub struct CLIPTextConfig {
    /// Vocabulary size
    pub vocab_size: usize,
    /// Hidden size
    pub hidden_size: usize,
    /// Number of layers
    pub num_layers: usize,
    /// Number of attention heads
    pub num_heads: usize,
    /// Maximum sequence length
    pub max_position_embeddings: usize,
}

impl Default for CLIPConfig {
    fn default() -> Self {
        CLIPConfig {
            embed_dim: 512,
            vision_config: CLIPVisionConfig::default(),
            text_config: CLIPTextConfig::default(),
            logit_scale_init_value: 2.6592, // ln(1/0.07)
        }
    }
}

impl Default for CLIPVisionConfig {
    fn default() -> Self {
        CLIPVisionConfig {
            image_size: 224,
            patch_size: 16,
            hidden_size: 768,
            num_layers: 12,
            num_heads: 12,
            mlp_ratio: 4,
        }
    }
}

impl Default for CLIPTextConfig {
    fn default() -> Self {
        CLIPTextConfig {
            vocab_size: 49408,
            hidden_size: 512,
            num_layers: 12,
            num_heads: 8,
            max_position_embeddings: 77,
        }
    }
}

impl CLIPConfig {
    /// CLIP ViT-B/32
    pub fn vit_b_32() -> Self {
        CLIPConfig {
            embed_dim: 512,
            vision_config: CLIPVisionConfig {
                image_size: 224,
                patch_size: 32,
                hidden_size: 768,
                num_layers: 12,
                num_heads: 12,
                mlp_ratio: 4,
            },
            text_config: CLIPTextConfig {
                vocab_size: 49408,
                hidden_size: 512,
                num_layers: 12,
                num_heads: 8,
                max_position_embeddings: 77,
            },
            logit_scale_init_value: 2.6592,
        }
    }
    
    /// CLIP ViT-B/16
    pub fn vit_b_16() -> Self {
        CLIPConfig {
            embed_dim: 512,
            vision_config: CLIPVisionConfig {
                image_size: 224,
                patch_size: 16,
                hidden_size: 768,
                num_layers: 12,
                num_heads: 12,
                mlp_ratio: 4,
            },
            text_config: CLIPTextConfig::default(),
            logit_scale_init_value: 2.6592,
        }
    }
    
    /// CLIP ViT-L/14
    pub fn vit_l_14() -> Self {
        CLIPConfig {
            embed_dim: 768,
            vision_config: CLIPVisionConfig {
                image_size: 224,
                patch_size: 14,
                hidden_size: 1024,
                num_layers: 24,
                num_heads: 16,
                mlp_ratio: 4,
            },
            text_config: CLIPTextConfig {
                vocab_size: 49408,
                hidden_size: 768,
                num_layers: 12,
                num_heads: 12,
                max_position_embeddings: 77,
            },
            logit_scale_init_value: 2.6592,
        }
    }
}

/// CLIP Vision Encoder (using Vision Transformer)
pub struct CLIPVisionEncoder {
    vit: VisionTransformer,
    projection: Linear,
}

impl CLIPVisionEncoder {
    /// Create new vision encoder
    pub fn new(config: &CLIPVisionConfig, embed_dim: usize) -> Self {
        // Convert to ViT config
        let vit_config = ViTConfig {
            image_size: config.image_size,
            patch_size: config.patch_size,
            in_channels: 3,
            embed_dim: config.hidden_size,
            num_layers: config.num_layers,
            num_heads: config.num_heads,
            mlp_dim: config.hidden_size * config.mlp_ratio,
            num_classes: 0, // No classification head
            dropout: 0.0,
        };
        
        let vit = VisionTransformer::new(vit_config);
        let projection = Linear::new(config.hidden_size, embed_dim);
        
        CLIPVisionEncoder { vit, projection }
    }
    
    /// Encode images
    pub fn forward(&self, images: &Tensor) -> Result<Tensor, String> {
        // Get ViT features (CLS token)
        let features = self.vit.forward(images)?;
        
        // Project to shared embedding space
        Ok(self.projection.forward(&features))
    }
}

/// CLIP Text Encoder
pub struct CLIPTextEncoder {
    token_embedding: Tensor,
    position_embedding: Tensor,
    layers: Vec<CLIPTextLayer>,
    ln_final: LayerNorm,
    projection: Linear,
}

impl CLIPTextEncoder {
    /// Create new text encoder
    pub fn new(config: &CLIPTextConfig, embed_dim: usize) -> Self {
        let token_embedding = Tensor::randn(&[config.vocab_size, config.hidden_size]);
        let position_embedding = Tensor::randn(&[config.max_position_embeddings, config.hidden_size]);
        
        let layers = (0..config.num_layers)
            .map(|_| CLIPTextLayer::new(config.hidden_size, config.num_heads))
            .collect();
        
        let ln_final = LayerNorm::new(config.hidden_size, 1e-5);
        let projection = Linear::new(config.hidden_size, embed_dim);
        
        CLIPTextEncoder {
            token_embedding,
            position_embedding,
            layers,
            ln_final,
            projection,
        }
    }
    
    /// Encode text
    pub fn forward(&self, input_ids: &Tensor) -> Result<Tensor, String> {
        let dims = input_ids.dims();
        let seq_length = dims[1];
        
        // Get token embeddings
        let mut hidden_states = self.get_token_embeddings(input_ids)?;
        
        // Add position embeddings
        hidden_states = self.add_position_embeddings(&hidden_states, seq_length)?;
        
        // Pass through transformer layers
        for layer in &self.layers {
            hidden_states = layer.forward(&hidden_states)?;
        }
        
        // Final layer norm
        hidden_states = self.ln_final.forward(&hidden_states)?;
        
        // Extract features at EOS token position (last token)
        let features = self.extract_eos_features(&hidden_states, seq_length)?;
        
        // Project to shared embedding space
        Ok(self.projection.forward(&features))
    }
    
    fn get_token_embeddings(&self, input_ids: &Tensor) -> Result<Tensor, String> {
        let ids_data = input_ids.data_f32();
        let embed_data = self.token_embedding.data_f32();
        let dims = input_ids.dims();
        let batch_size = dims[0];
        let seq_length = dims[1];
        let hidden_size = self.token_embedding.dims()[1];
        
        let mut result = Vec::with_capacity(batch_size * seq_length * hidden_size);
        
        for &id in ids_data.iter() {
            let idx = id as usize;
            let start = idx * hidden_size;
            let end = start + hidden_size;
            result.extend_from_slice(&embed_data[start..end]);
        }
        
        Tensor::from_slice(&result, &[batch_size, seq_length, hidden_size])
            .map_err(|e| format!("Failed to create embeddings: {:?}", e))
    }
    
    fn add_position_embeddings(&self, hidden_states: &Tensor, seq_length: usize) -> Result<Tensor, String> {
        let pos_embed_data = self.position_embedding.data_f32();
        let hidden_data = hidden_states.data_f32();
        let dims = hidden_states.dims();
        let hidden_size = dims[2];
        
        let mut result = Vec::with_capacity(hidden_data.len());
        
        for i in 0..hidden_data.len() {
            let pos = (i / hidden_size) % seq_length;
            let pos_idx = pos * hidden_size + (i % hidden_size);
            result.push(hidden_data[i] + pos_embed_data[pos_idx]);
        }
        
        Tensor::from_slice(&result, dims)
            .map_err(|e| format!("Failed to add position embeddings: {:?}", e))
    }
    
    fn extract_eos_features(&self, hidden_states: &Tensor, seq_length: usize) -> Result<Tensor, String> {
        let data = hidden_states.data_f32();
        let dims = hidden_states.dims();
        let batch_size = dims[0];
        let hidden_size = dims[2];
        
        let mut result = Vec::with_capacity(batch_size * hidden_size);
        
        // Extract last token for each batch
        for b in 0..batch_size {
            let start = (b * seq_length + seq_length - 1) * hidden_size;
            let end = start + hidden_size;
            result.extend_from_slice(&data[start..end]);
        }
        
        Tensor::from_slice(&result, &[batch_size, hidden_size])
            .map_err(|e| format!("Failed to extract EOS features: {:?}", e))
    }
}

/// CLIP Text Transformer Layer
pub struct CLIPTextLayer {
    self_attn: MultiHeadAttention,
    mlp: MLP,
    ln1: LayerNorm,
    ln2: LayerNorm,
}

impl CLIPTextLayer {
    fn new(hidden_size: usize, num_heads: usize) -> Self {
        CLIPTextLayer {
            self_attn: MultiHeadAttention::new(hidden_size, num_heads),
            mlp: MLP::new(hidden_size, hidden_size * 4),
            ln1: LayerNorm::new(hidden_size, 1e-5),
            ln2: LayerNorm::new(hidden_size, 1e-5),
        }
    }
    
    fn forward(&self, x: &Tensor) -> Result<Tensor, String> {
        // Self attention with residual
        let residual = x.clone();
        let x = self.ln1.forward(x)?;
        let x = self.self_attn.forward(&x)?;
        let x = x.add(&residual).unwrap_or(x);
        
        // MLP with residual
        let residual = x.clone();
        let x = self.ln2.forward(&x)?;
        let x = self.mlp.forward(&x)?;
        let x = x.add(&residual).unwrap_or(x);
        
        Ok(x)
    }
}

/// Multi-Head Attention (simplified)
pub struct MultiHeadAttention {
    q_proj: Linear,
    _k_proj: Linear,
    _v_proj: Linear,
    out_proj: Linear,
    _num_heads: usize,
    _head_dim: usize,
}

impl MultiHeadAttention {
    fn new(hidden_size: usize, num_heads: usize) -> Self {
        let head_dim = hidden_size / num_heads;
        MultiHeadAttention {
            q_proj: Linear::new(hidden_size, hidden_size),
            _k_proj: Linear::new(hidden_size, hidden_size),
            _v_proj: Linear::new(hidden_size, hidden_size),
            out_proj: Linear::new(hidden_size, hidden_size),
            _num_heads: num_heads,
            _head_dim: head_dim,
        }
    }
    
    fn forward(&self, x: &Tensor) -> Result<Tensor, String> {
        let q = self.q_proj.forward(x);
        // Simplified attention (real implementation would do proper multi-head)
        Ok(self.out_proj.forward(&q))
    }
}

/// MLP (Feed-Forward Network)
pub struct MLP {
    fc1: Linear,
    fc2: Linear,
}

impl MLP {
    fn new(hidden_size: usize, intermediate_size: usize) -> Self {
        MLP {
            fc1: Linear::new(hidden_size, intermediate_size),
            fc2: Linear::new(intermediate_size, hidden_size),
        }
    }
    
    fn forward(&self, x: &Tensor) -> Result<Tensor, String> {
        let x = self.fc1.forward(x);
        let x = x.gelu();
        Ok(self.fc2.forward(&x))
    }
}

/// Layer Normalization
pub struct LayerNorm {
    weight: Tensor,
    bias: Tensor,
    eps: f32,
}

impl LayerNorm {
    fn new(hidden_size: usize, eps: f32) -> Self {
        LayerNorm {
            weight: Tensor::ones(&[hidden_size]),
            bias: Tensor::zeros(&[hidden_size]),
            eps,
        }
    }
    
    fn forward(&self, x: &Tensor) -> Result<Tensor, String> {
        let x_data = x.data_f32();
        let dims = x.dims();
        let hidden_size = dims[dims.len() - 1];
        let batch_seq = x_data.len() / hidden_size;
        
        let weight_data = self.weight.data_f32();
        let bias_data = self.bias.data_f32();
        let mut result = Vec::with_capacity(x_data.len());
        
        for i in 0..batch_seq {
            let start = i * hidden_size;
            let end = start + hidden_size;
            let slice = &x_data[start..end];
            
            // Compute mean and variance
            let mean: f32 = slice.iter().sum::<f32>() / hidden_size as f32;
            let variance: f32 = slice.iter()
                .map(|x| (x - mean).powi(2))
                .sum::<f32>() / hidden_size as f32;
            let std = (variance + self.eps).sqrt();
            
            // Normalize and scale
            for (j, &x) in slice.iter().enumerate() {
                result.push((x - mean) / std * weight_data[j] + bias_data[j]);
            }
        }
        
        Tensor::from_slice(&result, dims)
            .map_err(|e| format!("Failed to normalize: {:?}", e))
    }
}

/// CLIP Model
pub struct CLIP {
    vision_encoder: CLIPVisionEncoder,
    text_encoder: CLIPTextEncoder,
    logit_scale: f32,
}

impl CLIP {
    /// Create new CLIP model
    pub fn new(config: CLIPConfig) -> Self {
        let vision_encoder = CLIPVisionEncoder::new(&config.vision_config, config.embed_dim);
        let text_encoder = CLIPTextEncoder::new(&config.text_config, config.embed_dim);
        let logit_scale = config.logit_scale_init_value.exp();
        
        CLIP {
            vision_encoder,
            text_encoder,
            logit_scale,
        }
    }
    
    /// Encode images
    pub fn encode_image(&self, images: &Tensor) -> Result<Tensor, String> {
        let features = self.vision_encoder.forward(images)?;
        Ok(self.normalize(&features))
    }
    
    /// Encode text
    pub fn encode_text(&self, input_ids: &Tensor) -> Result<Tensor, String> {
        let features = self.text_encoder.forward(input_ids)?;
        Ok(self.normalize(&features))
    }
    
    /// Forward pass (compute similarity matrix)
    pub fn forward(&self, images: &Tensor, input_ids: &Tensor) -> Result<Tensor, String> {
        let image_features = self.encode_image(images)?;
        let text_features = self.encode_text(input_ids)?;
        
        // Compute cosine similarity
        self.compute_similarity(&image_features, &text_features)
    }
    
    /// Normalize features (L2 normalization)
    fn normalize(&self, x: &Tensor) -> Tensor {
        let data = x.data_f32();
        let dims = x.dims();
        let feature_dim = dims[dims.len() - 1];
        let batch_size = data.len() / feature_dim;
        
        let mut result = Vec::with_capacity(data.len());
        
        for i in 0..batch_size {
            let start = i * feature_dim;
            let end = start + feature_dim;
            let slice = &data[start..end];
            
            // Compute L2 norm
            let norm: f32 = slice.iter().map(|x| x * x).sum::<f32>().sqrt();
            let norm = norm.max(1e-8); // Avoid division by zero
            
            // Normalize
            for &x in slice.iter() {
                result.push(x / norm);
            }
        }
        
        Tensor::from_slice(&result, dims).unwrap_or_else(|_| x.clone())
    }
    
    /// Compute similarity matrix
    fn compute_similarity(&self, image_features: &Tensor, text_features: &Tensor) -> Result<Tensor, String> {
        let img_data = image_features.data_f32();
        let txt_data = text_features.data_f32();
        
        let img_dims = image_features.dims();
        let txt_dims = text_features.dims();
        
        let num_images = img_dims[0];
        let num_texts = txt_dims[0];
        let feature_dim = img_dims[1];
        
        let mut result = Vec::with_capacity(num_images * num_texts);
        
        // Compute dot product (cosine similarity since features are normalized)
        for i in 0..num_images {
            for j in 0..num_texts {
                let mut dot_product = 0.0;
                for k in 0..feature_dim {
                    dot_product += img_data[i * feature_dim + k] * txt_data[j * feature_dim + k];
                }
                result.push(dot_product * self.logit_scale);
            }
        }
        
        Tensor::from_slice(&result, &[num_images, num_texts])
            .map_err(|e| format!("Failed to compute similarity: {:?}", e))
    }
    
    /// Zero-shot classification
    pub fn zero_shot_classify(&self, images: &Tensor, text_prompts: &Tensor) -> Result<Vec<usize>, String> {
        let similarity = self.forward(images, text_prompts)?;
        let data = similarity.data_f32();
        let dims = similarity.dims();
        let num_images = dims[0];
        let num_classes = dims[1];
        
        let mut predictions = Vec::with_capacity(num_images);
        
        for i in 0..num_images {
            let start = i * num_classes;
            let end = start + num_classes;
            let scores = &data[start..end];
            
            let pred = scores.iter()
                .enumerate()
                .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
                .map(|(idx, _)| idx)
                .unwrap_or(0);
            
            predictions.push(pred);
        }
        
        Ok(predictions)
    }
    
    /// Image-text retrieval (find best matching text for each image)
    pub fn image_to_text_retrieval(&self, images: &Tensor, texts: &Tensor) -> Result<Vec<usize>, String> {
        self.zero_shot_classify(images, texts)
    }
    
    /// Text-image retrieval (find best matching image for each text)
    pub fn text_to_image_retrieval(&self, images: &Tensor, texts: &Tensor) -> Result<Vec<usize>, String> {
        let similarity = self.forward(images, texts)?;
        let data = similarity.data_f32();
        let dims = similarity.dims();
        let num_images = dims[0];
        let num_texts = dims[1];
        
        let mut predictions = Vec::with_capacity(num_texts);
        
        // Transpose: for each text, find best image
        for j in 0..num_texts {
            let mut best_idx = 0;
            let mut best_score = data[j];
            
            for i in 1..num_images {
                let score = data[i * num_texts + j];
                if score > best_score {
                    best_score = score;
                    best_idx = i;
                }
            }
            
            predictions.push(best_idx);
        }
        
        Ok(predictions)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_clip_config() {
        let config = CLIPConfig::vit_b_32();
        assert_eq!(config.embed_dim, 512);
        assert_eq!(config.vision_config.patch_size, 32);
        
        let config = CLIPConfig::vit_l_14();
        assert_eq!(config.embed_dim, 768);
        assert_eq!(config.vision_config.num_layers, 24);
    }
    
    #[test]
    fn test_clip_vision_encoder() {
        let config = CLIPVisionConfig::default();
        let encoder = CLIPVisionEncoder::new(&config, 512);
        
        let images = Tensor::randn(&[2, 3, 224, 224]);
        let features = encoder.forward(&images).unwrap();
        
        assert_eq!(features.dims(), &[2, 512]);
    }
    
    #[test]
    fn test_clip_text_encoder() {
        let config = CLIPTextConfig::default();
        let encoder = CLIPTextEncoder::new(&config, 512);
        
        let input_ids = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]).unwrap();
        let features = encoder.forward(&input_ids).unwrap();
        
        assert_eq!(features.dims(), &[2, 512]);
    }
    
    #[test]
    fn test_clip_model() {
        let config = CLIPConfig::vit_b_32();
        let model = CLIP::new(config);
        
        let images = Tensor::randn(&[2, 3, 224, 224]);
        let input_ids = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]).unwrap();
        
        let similarity = model.forward(&images, &input_ids).unwrap();
        assert_eq!(similarity.dims(), &[2, 2]); // 2 images x 2 texts
    }
    
    #[test]
    fn test_zero_shot_classification() {
        let config = CLIPConfig::vit_b_32();
        let model = CLIP::new(config);
        
        let images = Tensor::randn(&[3, 3, 224, 224]);
        let text_prompts = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]).unwrap();
        
        let predictions = model.zero_shot_classify(&images, &text_prompts).unwrap();
        assert_eq!(predictions.len(), 3); // 3 images
    }
    
    #[test]
    fn test_layer_norm() {
        let ln = LayerNorm::new(128, 1e-5);
        let x = Tensor::randn(&[2, 4, 128]);
        let output = ln.forward(&x).unwrap();
        assert_eq!(output.dims(), &[2, 4, 128]);
    }
}