embeddings_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 EmbeddingsMatcher:
    """
    Uses Sentence Transformers to create embeddings for column names and values,
    then finds similarities between columns from two different CSV files.
    """
    
    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
        """
        Initialize the embeddings matcher with a pre-trained model
        
        Args:
            model_name: Name of the SentenceTransformer model to use
        """
        self.model_name = model_name
        try:
            self.model = SentenceTransformer(model_name)
            logger.info(f"Loaded SentenceTransformer 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', 'separate')
            
        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
        
        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
        # Note: cosine_similarity returns (n_samples_X, n_samples_Y) where each element [i,j] 
        # is the similarity between sample i of X and sample j of Y
        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
        
        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 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
    
    def match_columns_detailed(self, 
                              file1_columns: Dict[str, List[str]], 
                              file2_columns: Dict[str, List[str]],
                              threshold: float = 0.3) -> Dict[str, any]:
        """
        Detailed matching method that returns more information
        
        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:
            Dict containing matches, similarity matrix, and other details
        """
        # Compute similarity matrix
        similarity_matrix, file1_names, file2_names = self.compute_similarity_matrix(
            file1_columns, file2_columns
        )
        
        # Find all matches above threshold
        matches = self.find_best_matches(similarity_matrix, file1_names, file2_names, threshold)
        
        # Create detailed result
        result = {
            "matches": matches,
            "similarity_matrix": similarity_matrix,
            "file1_columns": file1_names,
            "file2_columns": file2_names,
            "details": {}
        }
        
        # Add top matches for each column
        for i, col1 in enumerate(file1_names):
            scores = similarity_matrix[i, :]
            top_idx = np.argmax(scores)
            result["details"][col1] = {
                "best_match": file2_names[top_idx],
                "best_score": float(scores[top_idx]),
                "all_scores": [(file2_names[j], float(scores[j])) for j in np.argsort(scores)[::-1]]
            }
        
        return result


def test_embeddings_matcher():
    """
    Test function for embeddings matcher
    """
    # 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"]
    }
    
    matcher = EmbeddingsMatcher()
    
    # Test basic matching
    matches = matcher.match_columns(file1_columns, file2_columns, threshold=0.3)
    print("Basic matches:")
    for col1, col2, score in matches:
        print(f"  {col1} ≈ {col2} (score: {score:.3f})")
    
    # Test detailed matching
    detailed_result = matcher.match_columns_detailed(file1_columns, file2_columns, threshold=0.3)
    print("\nDetailed matches:")
    for col1, details in detailed_result["details"].items():
        print(f"  {col1} -> {details['best_match']} (score: {details['best_score']:.3f})")
    
    print(f"\nFound {len(matches)} matches in total")


if __name__ == "__main__":
    test_embeddings_matcher()