fallback_matcher.py

from sentence_transformers import SentenceTransformer
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
from typing import List, Tuple, Dict, Optional
import logging

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class AlternativeEmbeddingsMatcher:
    """
    Uses lightweight models as alternatives to heavy models like Gemma.
    Uses efficient, smaller models that are under 300MB.
    """
    
    def __init__(self, model_name: str = "sentence-transformers/all-MiniLM-L6-v2"):
        """
        Initialize the alternative embeddings matcher with a lightweight model
        
        Args:
            model_name: Name of the lightweight model to use (under 300MB)
        """
        self.model_name = model_name
        try:
            logger.info(f"Loading lightweight model: {model_name}")
            self.model = SentenceTransformer(model_name)
            logger.info(f"Loaded lightweight model: {model_name}")
        except Exception as e:
            logger.error(f"Error loading model {model_name}: {str(e)}")
            raise
    
    def create_column_embedding(self, column_name: str, sample_values: List[str], 
                               combine_method: str = "concat") -> np.ndarray:
        """
        Create an embedding for a column based on its name and sample values
        
        Args:
            column_name: Name of the column
            sample_values: Sample values from the column
            combine_method: How to combine name and values ('concat', 'average')
            
        Returns:
            Embedding vector for the column
        """
        if combine_method == "concat":
            # Combine column name with sample values
            combined_text = f"{column_name} " + " ".join(sample_values[:5])  # Limit to 5 samples
            embeddings = self.model.encode([combined_text])
            return embeddings[0]  # Return single embedding vector
        
        elif combine_method == "average":
            # Create embeddings for name and values separately, then average
            name_embedding = self.model.encode([column_name])[0]
            if sample_values:
                # Take up to 5 sample values
                sample_text = " ".join(sample_values[:5])
                sample_embedding = self.model.encode([sample_text])[0]
                # Average the embeddings
                combined_embedding = (name_embedding + sample_embedding) / 2
                return combined_embedding
            else:
                return name_embedding
        else:  # separate
            # Return separate embeddings for name and values
            name_embedding = self.model.encode([column_name])[0]
            if sample_values:
                sample_text = " ".join(sample_values[:5])
                sample_embedding = self.model.encode([sample_text])[0]
                return np.concatenate([name_embedding, sample_embedding])
            else:
                # Pad with zeros to match dimensions when values are not available
                sample_embedding = np.zeros_like(name_embedding)
                return np.concatenate([name_embedding, sample_embedding])
    
    def compute_similarity_matrix(self, 
                                  file1_columns: Dict[str, List[str]], 
                                  file2_columns: Dict[str, List[str]]) -> Tuple[np.ndarray, List[str], List[str]]:
        """
        Compute similarity matrix between columns of two files using lightweight embeddings
        
        Args:
            file1_columns: Dict mapping column names to sample values for file 1
            file2_columns: Dict mapping column names to sample values for file 2
            
        Returns:
            Tuple of (similarity_matrix, file1_column_names, file2_column_names)
        """
        # Get column names
        file1_names = list(file1_columns.keys())
        file2_names = list(file2_columns.keys())
        
        # Create embeddings for all columns in file 1
        file1_embeddings = []
        for col_name in file1_names:
            embedding = self.create_column_embedding(col_name, file1_columns[col_name])
            file1_embeddings.append(embedding)
        
        # Create embeddings for all columns in file 2
        file2_embeddings = []
        for col_name in file2_names:
            embedding = self.create_column_embedding(col_name, file2_columns[col_name])
            file2_embeddings.append(embedding)
        
        # Convert to numpy arrays
        file1_embeddings = np.array(file1_embeddings)  # Shape: (n_cols1, embedding_dim)
        file2_embeddings = np.array(file2_embeddings)  # Shape: (n_cols2, embedding_dim)
        
        # Compute cosine similarity matrix
        similarity_matrix = cosine_similarity(file1_embeddings, file2_embeddings)
        
        return similarity_matrix, file1_names, file2_names
    
    def find_best_matches(self, 
                          similarity_matrix: np.ndarray, 
                          file1_names: List[str], 
                          file2_names: List[str],
                          threshold: float = 0.5) -> List[Tuple[str, str, float]]:
        """
        Find the best matching column pairs above a similarity threshold
        
        Args:
            similarity_matrix: Matrix of similarities between file1 and file2 columns
            file1_names: Names of columns in file 1
            file2_names: Names of columns in file 2
            threshold: Minimum similarity score to consider a match
            
        Returns:
            List of tuples (file1_col, file2_col, similarity_score) for matches above threshold
        """
        matches = []
        
        # For each column in file 1, find the best matching column in file 2
        for i, col1 in enumerate(file1_names):
            best_match_idx = np.argmax(similarity_matrix[i, :])
            best_score = similarity_matrix[i, best_match_idx]
            col2 = file2_names[best_match_idx]
            
            if best_score >= threshold:
                matches.append((col1, col2, float(best_score)))
        
        # Also check for each column in file 2, find the best matching column in file 1
        # This helps find matches that might be missed in the previous step
        for j, col2 in enumerate(file2_names):
            best_match_idx = np.argmax(similarity_matrix[:, j])
            best_score = similarity_matrix[best_match_idx, j]
            col1 = file1_names[best_match_idx]
            
            # Check if this match is already in the list
            match_exists = any(m[0] == col1 and m[1] == col2 for m in matches)
            if not match_exists and best_score >= threshold:
                matches.append((col1, col2, float(best_score)))
        
        # Remove duplicates and sort by similarity score in descending order
        matches = list(set(matches))
        matches.sort(key=lambda x: x[2], reverse=True)
        
        return matches
    
    def match_columns(self, 
                      file1_columns: Dict[str, List[str]], 
                      file2_columns: Dict[str, List[str]],
                      threshold: float = 0.5) -> List[Tuple[str, str, float]]:
        """
        Main method to match columns between two files using lightweight embeddings
        
        Args:
            file1_columns: Dict mapping column names to sample values for file 1
            file2_columns: Dict mapping column names to sample values for file 2
            threshold: Minimum similarity score to consider a match
            
        Returns:
            List of tuples (file1_col, file2_col, similarity_score) for matches above threshold
        """
        logger.info(f"Matching columns using lightweight model {self.model_name} between two files with {len(file1_columns)} and {len(file2_columns)} columns")
        
        # Compute similarity matrix
        similarity_matrix, file1_names, file2_names = self.compute_similarity_matrix(
            file1_columns, file2_columns
        )
        
        # Find best matches
        matches = self.find_best_matches(similarity_matrix, file1_names, file2_names, threshold)
        
        logger.info(f"Found {len(matches)} matches above threshold {threshold}")
        
        return matches


class FallbackEmbeddingsMatcher:
    """
    A wrapper class that tries multiple lightweight models as fallbacks
    when the primary model doesn't find enough matches.
    """
    
    def __init__(self, 
                 primary_model: str = "sentence-transformers/all-MiniLM-L6-v2",
                 fallback_models: List[str] = [
                     "sentence-transformers/paraphrase-MiniLM-L6-v2",
                     "sentence-transformers/all-MiniLM-L12-v2",
                     "nomic-ai/nomic-embed-text-v1.5"
                 ]):
        """
        Initialize with a primary model and multiple fallback options
        
        Args:
            primary_model: Primary lightweight model (under 300MB)
            fallback_models: List of fallback models to try if primary doesn't find enough matches
        """
        self.primary_model_name = primary_model
        self.fallback_models = fallback_models
        
        # Try to initialize primary model
        try:
            self.primary_model = SentenceTransformer(primary_model)
            self.primary_available = True
            logger.info(f"Initialized primary model: {primary_model}")
        except Exception as e:
            logger.warning(f"Could not load primary model {primary_model}: {e}")
            self.primary_available = False
            self.primary_model = None
    
    def match_columns(self, 
                      file1_columns: Dict[str, List[str]], 
                      file2_columns: Dict[str, List[str]],
                      threshold: float = 0.5,
                      min_matches_for_primary: int = 2,
                      use_fallback_when_unmatched: bool = True) -> List[Tuple[str, str, float]]:
        """
        Match columns using primary method with fallback option
        
        Args:
            file1_columns: Dict mapping column names to sample values for file 1
            file2_columns: Dict mapping column names to sample values for file 2
            threshold: Minimum similarity score to consider a match
            min_matches_for_primary: Minimum number of matches needed to skip fallback
            use_fallback_when_unmatched: Whether to use fallback for unmatched columns
            
        Returns:
            List of tuples (file1_col, file2_col, similarity_score) for matches above threshold
        """
        
        # Start with primary model
        all_matches = []
        unmatched_file1 = set(file1_columns.keys())
        unmatched_file2 = set(file2_columns.keys())
        
        if self.primary_available:
            try:
                from embeddings_matcher import EmbeddingsMatcher
                # Use the primary model
                primary_matcher = EmbeddingsMatcher(self.primary_model_name)
                primary_matches = primary_matcher.match_columns(
                    file1_columns, 
                    file2_columns, 
                    threshold
                )
                
                all_matches.extend(primary_matches)
                
                # Update unmatched columns based on primary matches
                for col1, col2, _ in primary_matches:
                    unmatched_file1.discard(col1)
                    unmatched_file2.discard(col2)
                
                logger.info(f"Primary model found {len(primary_matches)} matches")
                
                # If we don't have many matches, try fallback models for unmatched columns
                if len(primary_matches) < min_matches_for_primary and use_fallback_when_unmatched:
                    logger.info(f"Primary model found few matches, trying fallback on unmatched columns")
                    
                    # Create subsets of unmatched columns
                    unmatched_file1_info = {col: file1_columns[col] for col in unmatched_file1}
                    unmatched_file2_info = {col: file2_columns[col] for col in unmatched_file2}
                    
                    # Try fallback models for unmatched columns
                    for fallback_model in self.fallback_models:
                        try:
                            logger.info(f"Trying fallback model: {fallback_model}")
                            fallback_matcher = AlternativeEmbeddingsMatcher(fallback_model)
                            fallback_matches = fallback_matcher.match_columns(
                                unmatched_file1_info,
                                unmatched_file2_info,
                                threshold
                            )
                            
                            # Add fallback matches and update unmatched sets
                            for col1, col2, score in fallback_matches:
                                if col1 in unmatched_file1 and col2 in unmatched_file2:
                                    all_matches.append((col1, col2, score))
                                    unmatched_file1.discard(col1)
                                    unmatched_file2.discard(col2)
                            
                            # If we found good matches with this fallback, stop
                            if len(fallback_matches) > 0:
                                logger.info(f"Fallback model {fallback_model} found {len(fallback_matches)} matches")
                                break
                            
                        except Exception as e:
                            logger.warning(f"Fallback model {fallback_model} failed: {e}")
                            continue
                            
            except Exception as e:
                logger.error(f"Primary model failed: {e}")
        
        # If no primary model is available, try fallback models
        elif use_fallback_when_unmatched:
            logger.info("Primary model not available, trying fallback models")
            for fallback_model in self.fallback_models:
                try:
                    logger.info(f"Trying fallback model: {fallback_model}")
                    fallback_matcher = AlternativeEmbeddingsMatcher(fallback_model)
                    fallback_matches = fallback_matcher.match_columns(
                        file1_columns,
                        file2_columns,
                        threshold
                    )
                    
                    if fallback_matches:
                        logger.info(f"Fallback model {fallback_model} found {len(fallback_matches)} matches")
                        return fallback_matches
                        
                except Exception as e:
                    logger.warning(f"Fallback model {fallback_model} failed: {e}")
                    continue
        
        # Remove duplicates and sort by similarity score in descending order
        all_matches = list(set(all_matches))
        all_matches.sort(key=lambda x: x[2], reverse=True)
        
        logger.info(f"Total matches found: {len(all_matches)}")
        logger.info(f"Unmatched in file1: {len(unmatched_file1)}, Unmatched in file2: {len(unmatched_file2)}")
        
        return all_matches


def test_alternative_matcher():
    """
    Test function for alternative embeddings matcher using lightweight models
    """
    # Create sample data
    file1_columns = {
        "Project_ID": ["1", "2", "3"],
        "Cost_Center": ["CC101", "CC102", "CC103"],
        "Expense_Total": ["15000", "25000", "30000"]
    }
    
    file2_columns = {
        "Engagement_ID": ["101", "102", "103"],
        "Org_Code": ["CC101", "CC102", "CC103"],
        "Actual_Spend": ["14800", "25200", "29500"]
    }
    
    # Test with different lightweight models
    models_to_test = [
        "sentence-transformers/all-MiniLM-L6-v2",
        "sentence-transformers/paraphrase-MiniLM-L6-v2",
        "sentence-transformers/all-MiniLM-L12-v2"
    ]
    
    for model_name in models_to_test:
        try:
            logger.info(f"Testing model: {model_name}")
            matcher = AlternativeEmbeddingsMatcher(model_name)
            matches = matcher.match_columns(file1_columns, file2_columns, threshold=0.3)
            print(f"\n{model_name} matches:")
            for col1, col2, score in matches:
                print(f"  {col1} ≈ {col2} (score: {score:.3f})")
        except Exception as e:
            print(f"Model {model_name} test failed: {e}")


def test_fallback_matcher():
    """
    Test the enhanced fallback mechanism
    """
    # Create sample data
    file1_columns = {
        "Project_ID": ["1", "2", "3"],
        "Cost_Center": ["CC101", "CC102", "CC103"],
        "Expense_Total": ["15000", "25000", "30000"]
    }
    
    file2_columns = {
        "Engagement_ID": ["101", "102", "103"],
        "Org_Code": ["CC101", "CC102", "CC103"],
        "Actual_Spend": ["14800", "25200", "29500"]
    }
    
    try:
        # Create fallback matcher with multiple options
        fallback_models = [
            "sentence-transformers/paraphrase-MiniLM-L6-v2",
            "sentence-transformers/all-MiniLM-L12-v2",
            "nomic-ai/nomic-embed-text-v1.5"
        ]
        
        fallback_matcher = FallbackEmbeddingsMatcher(
            primary_model="sentence-transformers/all-MiniLM-L6-v2",
            fallback_models=fallback_models
        )
        
        matches = fallback_matcher.match_columns(file1_columns, file2_columns, threshold=0.3)
        print("\nFallback matcher results:")
        for col1, col2, score in matches:
            print(f"  {col1} ≈ {col2} (score: {score:.3f})")
        
        print(f"\nTotal matches found: {len(matches)}")
        
    except Exception as e:
        print(f"Fallback matcher test failed: {e}")


if __name__ == "__main__":
    test_alternative_matcher()
    print("\n" + "="*50 + "\n")
    test_fallback_matcher()