cross_modal_fusion.py

#!/usr/bin/env python3
"""
Cross-Modal Fusion Engine - Unique Feature for Advanced Multimodal Intelligence
Handles sophisticated fusion of embeddings across text, image, and audio modalities
"""

import numpy as np
import torch
import torch.nn as nn
from typing import Dict, Any, List, Optional, Tuple
import logging
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.decomposition import PCA
import faiss

logger = logging.getLogger(__name__)

class CrossModalFusionEngine:
    """Advanced cross-modal embedding fusion and intelligence system"""
    
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # Embedding dimensions for different modalities
        self.text_dim = 384  # all-MiniLM-L6-v2 dimension
        self.image_dim = 512  # CLIP dimension
        self.audio_dim = 384  # Same as text after transcription
        
        # Fusion parameters
        self.fusion_dim = 512  # Target fusion dimension
        self.similarity_threshold = config['search']['similarity_threshold']
        
        # Initialize fusion networks
        self._initialize_fusion_networks()
        
        # Cross-modal mapping matrices
        self.cross_modal_mappings = {}
        
        # Semantic relationship tracking
        self.semantic_clusters = {}
        self.modality_bridges = {}
        
    def _initialize_fusion_networks(self):
        """Initialize neural networks for cross-modal fusion"""
        
        # Text-to-unified projection
        self.text_projector = nn.Sequential(
            nn.Linear(self.text_dim, self.fusion_dim),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(self.fusion_dim, self.fusion_dim),
            nn.LayerNorm(self.fusion_dim)
        ).to(self.device)
        
        # Image-to-unified projection
        self.image_projector = nn.Sequential(
            nn.Linear(self.image_dim, self.fusion_dim),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(self.fusion_dim, self.fusion_dim),
            nn.LayerNorm(self.fusion_dim)
        ).to(self.device)
        
        # Audio-to-unified projection (after transcription)
        self.audio_projector = nn.Sequential(
            nn.Linear(self.audio_dim, self.fusion_dim),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(self.fusion_dim, self.fusion_dim),
            nn.LayerNorm(self.fusion_dim)
        ).to(self.device)
        
        # Attention-based fusion mechanism
        self.fusion_attention = nn.MultiheadAttention(
            embed_dim=self.fusion_dim,
            num_heads=8,
            dropout=0.1,
            batch_first=True
        ).to(self.device)
        
        # Final fusion layer
        self.fusion_combiner = nn.Sequential(
            nn.Linear(self.fusion_dim * 3, self.fusion_dim),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(self.fusion_dim, self.fusion_dim),
            nn.LayerNorm(self.fusion_dim)
        ).to(self.device)
        
    def create_unified_embedding(self, 
                               text_embedding: Optional[np.ndarray] = None,
                               image_embedding: Optional[np.ndarray] = None,
                               audio_embedding: Optional[np.ndarray] = None,
                               fusion_weights: Optional[Dict[str, float]] = None) -> np.ndarray:
        """Create unified cross-modal embedding"""
        
        fusion_weights = fusion_weights or {'text': 1.0, 'image': 1.0, 'audio': 1.0}
        
        with torch.no_grad():
            embeddings = []
            weights = []
            
            # Process text embedding
            if text_embedding is not None:
                text_tensor = torch.from_numpy(text_embedding).float().to(self.device)
                if text_tensor.dim() == 1:
                    text_tensor = text_tensor.unsqueeze(0)
                text_unified = self.text_projector(text_tensor)
                embeddings.append(text_unified)
                weights.append(fusion_weights['text'])
                
            # Process image embedding
            if image_embedding is not None:
                image_tensor = torch.from_numpy(image_embedding).float().to(self.device)
                if image_tensor.dim() == 1:
                    image_tensor = image_tensor.unsqueeze(0)
                image_unified = self.image_projector(image_tensor)
                embeddings.append(image_unified)
                weights.append(fusion_weights['image'])
                
            # Process audio embedding
            if audio_embedding is not None:
                audio_tensor = torch.from_numpy(audio_embedding).float().to(self.device)
                if audio_tensor.dim() == 1:
                    audio_tensor = audio_tensor.unsqueeze(0)
                audio_unified = self.audio_projector(audio_tensor)
                embeddings.append(audio_unified)
                weights.append(fusion_weights['audio'])
                
            if not embeddings:
                raise ValueError("At least one embedding must be provided")
                
            # Weighted combination
            weights_tensor = torch.tensor(weights).float().to(self.device)
            weights_tensor = weights_tensor / weights_tensor.sum()  # Normalize
            
            # Stack embeddings and apply weights
            stacked_embeddings = torch.stack(embeddings, dim=1)  # [batch, num_modalities, fusion_dim]
            
            # Apply attention-based fusion
            attended_embeddings, attention_weights = self.fusion_attention(
                stacked_embeddings, stacked_embeddings, stacked_embeddings
            )
            
            # Weighted average with learned attention
            weighted_embeddings = (attended_embeddings * weights_tensor.view(1, -1, 1)).sum(dim=1)
            
            return weighted_embeddings.cpu().numpy().flatten()
            
    def find_cross_modal_relationships(self, 
                                     embeddings_db: Dict[str, Dict[str, Any]],
                                     relationship_threshold: float = 0.7) -> Dict[str, List[Dict[str, Any]]]:
        """Find semantic relationships across different modalities"""
        
        relationships = {
            'text_image': [],
            'text_audio': [],
            'image_audio': [],
            'clusters': []
        }
        
        # Group embeddings by modality
        modality_groups = {'text': [], 'image': [], 'audio_transcription': []}
        
        for doc_id, doc_data in embeddings_db.items():
            modality = doc_data.get('modality', 'unknown')
            if modality in modality_groups:
                modality_groups[modality].append({
                    'id': doc_id,
                    'embedding': np.array(doc_data.get('embedding', [])),
                    'content': doc_data.get('content', ''),
                    'metadata': doc_data
                })
                
        # Find cross-modal similarities
        for mod1, mod2 in [('text', 'image'), ('text', 'audio_transcription'), ('image', 'audio_transcription')]:
            if not modality_groups[mod1] or not modality_groups[mod2]:
                continue
                
            # Create unified embeddings for comparison
            group1_unified = []
            group2_unified = []
            
            for item1 in modality_groups[mod1]:
                if mod1 == 'text':
                    unified = self.create_unified_embedding(text_embedding=item1['embedding'])
                elif mod1 == 'image':
                    unified = self.create_unified_embedding(image_embedding=item1['embedding'])
                else:  # audio
                    unified = self.create_unified_embedding(audio_embedding=item1['embedding'])
                group1_unified.append(unified)
                
            for item2 in modality_groups[mod2]:
                if mod2 == 'text':
                    unified = self.create_unified_embedding(text_embedding=item2['embedding'])
                elif mod2 == 'image':
                    unified = self.create_unified_embedding(image_embedding=item2['embedding'])
                else:  # audio
                    unified = self.create_unified_embedding(audio_embedding=item2['embedding'])
                group2_unified.append(unified)
                
            # Calculate similarities
            similarity_matrix = cosine_similarity(group1_unified, group2_unified)
            
            # Find high-similarity pairs
            high_sim_pairs = np.where(similarity_matrix > relationship_threshold)
            
            relationship_key = f"{mod1}_{mod2}"
            for i, j in zip(high_sim_pairs[0], high_sim_pairs[1]):
                relationship = {
                    'source_id': modality_groups[mod1][i]['id'],
                    'target_id': modality_groups[mod2][j]['id'],
                    'source_modality': mod1,
                    'target_modality': mod2,
                    'similarity_score': float(similarity_matrix[i, j]),
                    'source_content': modality_groups[mod1][i]['content'][:100],
                    'target_content': modality_groups[mod2][j]['content'][:100],
                    'relationship_type': self._classify_relationship_type(
                        modality_groups[mod1][i], modality_groups[mod2][j]
                    )
                }
                relationships[relationship_key].append(relationship)
                
        # Perform semantic clustering
        relationships['clusters'] = self._perform_semantic_clustering(embeddings_db)
        
        return relationships
        
    def _classify_relationship_type(self, item1: Dict[str, Any], item2: Dict[str, Any]) -> str:
        """Classify the type of relationship between two items"""
        
        # Simple heuristic classification - can be enhanced with ML
        content1 = item1['content'].lower()
        content2 = item2['content'].lower()
        
        # Check for direct content overlap
        words1 = set(content1.split())
        words2 = set(content2.split())
        overlap = len(words1.intersection(words2))
        
        if overlap > 5:
            return "content_overlap"
        elif overlap > 2:
            return "semantic_similarity"
        else:
            return "conceptual_relationship"
            
    def _perform_semantic_clustering(self, 
                                   embeddings_db: Dict[str, Dict[str, Any]], 
                                   n_clusters: int = 10) -> List[Dict[str, Any]]:
        """Perform semantic clustering across all modalities"""
        
        try:
            from sklearn.cluster import KMeans
            
            # Collect all embeddings and create unified representations
            all_embeddings = []
            all_metadata = []
            
            for doc_id, doc_data in embeddings_db.items():
                embedding = np.array(doc_data.get('embedding', []))
                modality = doc_data.get('modality', 'unknown')
                
                if len(embedding) == 0:
                    continue
                    
                # Create unified embedding based on modality
                if modality == 'text':
                    unified = self.create_unified_embedding(text_embedding=embedding)
                elif modality == 'image':
                    unified = self.create_unified_embedding(image_embedding=embedding)
                elif modality == 'audio_transcription':
                    unified = self.create_unified_embedding(audio_embedding=embedding)
                else:
                    continue
                    
                all_embeddings.append(unified)
                all_metadata.append({
                    'id': doc_id,
                    'modality': modality,
                    'content': doc_data.get('content', ''),
                    'file_name': doc_data.get('file_name', 'unknown')
                })
                
            if len(all_embeddings) < n_clusters:
                n_clusters = max(1, len(all_embeddings) // 2)
                
            # Perform clustering
            embeddings_matrix = np.vstack(all_embeddings)
            kmeans = KMeans(n_clusters=n_clusters, random_state=42)
            cluster_labels = kmeans.fit_predict(embeddings_matrix)
            
            # Organize clusters
            clusters = []
            for cluster_id in range(n_clusters):
                cluster_indices = np.where(cluster_labels == cluster_id)[0]
                
                if len(cluster_indices) == 0:
                    continue
                    
                cluster_items = [all_metadata[i] for i in cluster_indices]
                cluster_centroid = kmeans.cluster_centers_[cluster_id]
                
                # Analyze cluster composition
                modality_counts = {}
                for item in cluster_items:
                    mod = item['modality']
                    modality_counts[mod] = modality_counts.get(mod, 0) + 1
                    
                # Find representative item (closest to centroid)
                cluster_embeddings = embeddings_matrix[cluster_indices]
                distances = np.linalg.norm(cluster_embeddings - cluster_centroid, axis=1)
                representative_idx = cluster_indices[np.argmin(distances)]
                
                cluster_info = {
                    'cluster_id': cluster_id,
                    'size': len(cluster_items),
                    'modality_distribution': modality_counts,
                    'representative_item': all_metadata[representative_idx],
                    'items': cluster_items,
                    'centroid': cluster_centroid.tolist(),
                    'intra_cluster_similarity': float(np.mean(1 - distances / np.max(distances)))
                }
                
                clusters.append(cluster_info)
                
            return sorted(clusters, key=lambda x: x['size'], reverse=True)
            
        except Exception as e:
            logger.error(f"Semantic clustering failed: {e}")
            return []
            
    def enhance_search_with_fusion(self, 
                                 query_embeddings: Dict[str, np.ndarray],
                                 candidate_embeddings: List[Dict[str, Any]],
                                 fusion_strategy: str = "weighted_average") -> List[Dict[str, Any]]:
        """Enhance search results using cross-modal fusion"""
        
        enhanced_results = []
        
        # Create unified query embedding
        query_unified = self.create_unified_embedding(
            text_embedding=query_embeddings.get('text'),
            image_embedding=query_embeddings.get('image'),
            audio_embedding=query_embeddings.get('audio')
        )
        
        for candidate in candidate_embeddings:
            candidate_embedding = np.array(candidate.get('embedding', []))
            candidate_modality = candidate.get('modality', 'unknown')
            
            if len(candidate_embedding) == 0:
                continue
                
            # Create unified candidate embedding
            if candidate_modality == 'text':
                candidate_unified = self.create_unified_embedding(text_embedding=candidate_embedding)
            elif candidate_modality == 'image':
                candidate_unified = self.create_unified_embedding(image_embedding=candidate_embedding)
            elif candidate_modality == 'audio_transcription':
                candidate_unified = self.create_unified_embedding(audio_embedding=candidate_embedding)
            else:
                continue
                
            # Calculate enhanced similarity
            fusion_similarity = cosine_similarity(
                query_unified.reshape(1, -1),
                candidate_unified.reshape(1, -1)
            )[0, 0]
            
            # Combine with original similarity
            original_similarity = candidate.get('similarity_score', 0.0)
            
            if fusion_strategy == "weighted_average":
                enhanced_score = 0.6 * fusion_similarity + 0.4 * original_similarity
            elif fusion_strategy == "max":
                enhanced_score = max(fusion_similarity, original_similarity)
            else:  # default to fusion-only
                enhanced_score = fusion_similarity
                
            # Add fusion metadata
            enhanced_candidate = candidate.copy()
            enhanced_candidate.update({
                'fusion_similarity': float(fusion_similarity),
                'original_similarity': original_similarity,
                'enhanced_similarity': float(enhanced_score),
                'fusion_strategy': fusion_strategy
            })
            
            enhanced_results.append(enhanced_candidate)
            
        # Sort by enhanced similarity
        enhanced_results.sort(key=lambda x: x['enhanced_similarity'], reverse=True)
        
        return enhanced_results
        
    def generate_cross_modal_insights(self, 
                                    search_results: List[Dict[str, Any]]) -> Dict[str, Any]:
        """Generate insights about cross-modal patterns in search results"""
        
        insights = {
            'modality_distribution': {},
            'cross_modal_clusters': [],
            'semantic_bridges': [],
            'content_overlap_analysis': {},
            'fusion_effectiveness': {}
        }
        
        # Analyze modality distribution
        for result in search_results:
            modality = result.get('modality', 'unknown')
            insights['modality_distribution'][modality] = insights['modality_distribution'].get(modality, 0) + 1
            
        # Identify semantic bridges (items that connect different modalities)
        modality_groups = {}
        for result in search_results:
            modality = result.get('modality', 'unknown')
            if modality not in modality_groups:
                modality_groups[modality] = []
            modality_groups[modality].append(result)
            
        # Find items with high cross-modal similarity
        for mod1 in modality_groups:
            for mod2 in modality_groups:
                if mod1 != mod2:
                    # This is a simplified analysis - could be more sophisticated
                    bridge_candidates = []
                    
                    for item1 in modality_groups[mod1][:3]:  # Top 3 from each modality
                        for item2 in modality_groups[mod2][:3]:
                            # Check content similarity or other metrics
                            content_sim = self._calculate_content_similarity(
                                item1.get('content', ''),
                                item2.get('content', '')
                            )
                            
                            if content_sim > 0.3:
                                bridge_candidates.append({
                                    'item1': item1['id'],
                                    'item2': item2['id'],
                                    'modalities': [mod1, mod2],
                                    'content_similarity': content_sim
                                })
                                
                    if bridge_candidates:
                        insights['semantic_bridges'].extend(bridge_candidates)
                        
        return insights
        
    def _calculate_content_similarity(self, content1: str, content2: str) -> float:
        """Calculate simple content similarity between two texts"""
        
        words1 = set(content1.lower().split())
        words2 = set(content2.lower().split())
        
        if not words1 or not words2:
            return 0.0
            
        intersection = words1.intersection(words2)
        union = words1.union(words2)
        
        return len(intersection) / len(union) if union else 0.0