refactor: computeSimilarities

README.md

@@ -54,8 +54,24 @@ docker run -p 7860:7860 imagebind-api
 
 The API will be available at `http://localhost:7860` with the following endpoints:
 
-- POST `/compute_embeddings`: Generate embeddings for images, audio files, and text
-- POST `/compute_similarities`: Compute similarities between embeddings
+### POST `/compute_embeddings`
+
+Generate embeddings for images, audio files, and text.
+
+### POST `/compute_similarities`
+
+Compute similarities between embeddings with advanced filtering options:
+
+- Threshold filtering for minimum similarity scores
+- Top-K results limitation
+- Optional self-similarity inclusion
+- Score normalization
+- Detailed match information including original file/text references
+- Statistical analysis of similarity scores
+
+### GET `/health`
+
+Basic health check endpoint
 
 For detailed API documentation, visit `http://localhost:7860/docs`
 

main.py

@@ -119,8 +119,62 @@ class EmbeddingResponse(BaseModel):
     embeddings: dict
     file_names: dict
 
+class SimilarityRequest(BaseModel):
+    embeddings: Dict[str, List[List[float]]]
+    threshold: float = 0.5
+    top_k: int | None = None
+    include_self_similarity: bool = False
+    normalize_scores: bool = True
+
+class SimilarityMatch(BaseModel):
+    index_a: int
+    index_b: int
+    score: float
+    modality_a: str
+    modality_b: str
+    item_a: str  # Original item identifier (filename or text)
+    item_b: str  # Original item identifier (filename or text)
+
 class SimilarityResponse(BaseModel):
-    similarities: dict
+    matches: List[SimilarityMatch]
+    statistics: Dict[str, float]  # Contains avg_score, max_score, etc.
+    modality_pairs: List[str]  # Lists which modality comparisons were performed
+
+class ModalityPair:
+    def __init__(self, mod1: str, mod2: str):
+        self.mod1 = min(mod1, mod2)  # Ensure consistent ordering
+        self.mod2 = max(mod1, mod2)
+    
+    def __str__(self):
+        return f"{self.mod1}_to_{self.mod2}"
+
+def compute_similarity_matrix(tensor1: torch.Tensor, tensor2: torch.Tensor, normalize: bool = True) -> torch.Tensor:
+    """Compute cosine similarity between two sets of embeddings."""
+    # Normalize embeddings if requested
+    if normalize:
+        tensor1 = torch.nn.functional.normalize(tensor1, dim=1)
+        tensor2 = torch.nn.functional.normalize(tensor2, dim=1)
+    
+    # Compute similarity matrix
+    similarity = torch.matmul(tensor1, tensor2.T)
+    
+    return similarity
+
+def get_top_k_matches(similarity_matrix: torch.Tensor, top_k: int | None = None) -> List[tuple]:
+    """Get top-k matches from a similarity matrix."""
+    if top_k is None:
+        top_k = similarity_matrix.numel()
+    
+    # Flatten and get top-k indices
+    flat_sim = similarity_matrix.flatten()
+    top_k = min(top_k, flat_sim.numel())
+    values, indices = torch.topk(flat_sim, k=top_k)
+    
+    # Convert flat indices to 2D indices
+    rows = indices // similarity_matrix.size(1)
+    cols = indices % similarity_matrix.size(1)
+    
+    return [(r.item(), c.item(), v.item()) for r, c, v in zip(rows, cols, values)]
 
 @app.post("/compute_embeddings", response_model=EmbeddingResponse)
 async def generate_embeddings(
@@ -202,12 +256,96 @@ async def generate_embeddings(
 
 @app.post("/compute_similarities", response_model=SimilarityResponse)
 async def compute_similarities(
-    embeddings: Dict[str, List[List[float]]],
+    request: SimilarityRequest,
+    file_names: Dict[str, List[str]],  # Maps modality to list of file/text names
     credentials: HTTPAuthorizationCredentials = Depends(verify_token)
 ):
-    """Compute similarities from provided embeddings."""
-    similarities = embedding_manager.compute_similarities(embeddings)
-    return SimilarityResponse(similarities=similarities)
+    """
+    Compute cross-modal similarities with advanced filtering and matching options.
+    
+    Parameters:
+    - embeddings: Dict mapping modality to embedding tensors
+    - threshold: Minimum similarity score to include in results
+    - top_k: Maximum number of matches to return (per modality pair)
+    - include_self_similarity: Whether to include same-item comparisons
+    - normalize_scores: Whether to normalize embeddings before comparison
+    - file_names: Dict mapping modality to list of original file/text names
+    """
+    
+    matches = []
+    statistics = {
+        "avg_score": 0.0,
+        "max_score": 0.0,
+        "min_score": 1.0,
+        "total_comparisons": 0
+    }
+    
+    # Convert embeddings to tensors
+    tensors = {
+        k: torch.tensor(v) for k, v in request.embeddings.items()
+        if isinstance(v, (list, np.ndarray)) and len(v) > 0
+    }
+    
+    modality_pairs = []
+    all_scores = []
+    
+    # Get all possible modality pairs
+    modalities = list(tensors.keys())
+    for i, mod1 in enumerate(modalities):
+        for mod2 in modalities[i:]:  # Include self-comparisons if requested
+            if mod1 == mod2 and not request.include_self_similarity:
+                continue
+                
+            pair = ModalityPair(mod1, mod2)
+            modality_pairs.append(str(pair))
+            
+            # Compute similarity matrix
+            sim_matrix = compute_similarity_matrix(
+                tensors[mod1], 
+                tensors[mod2],
+                normalize=request.normalize_scores
+            )
+            
+            # Get top matches
+            top_matches = get_top_k_matches(sim_matrix, request.top_k)
+            
+            # Filter by threshold and create match objects
+            for idx_a, idx_b, score in top_matches:
+                if score < request.threshold:
+                    continue
+                    
+                # Skip self-matches if not requested
+                if mod1 == mod2 and idx_a == idx_b and not request.include_self_similarity:
+                    continue
+                
+                matches.append(SimilarityMatch(
+                    index_a=idx_a,
+                    index_b=idx_b,
+                    score=float(score),
+                    modality_a=mod1,
+                    modality_b=mod2,
+                    item_a=file_names[mod1][idx_a],
+                    item_b=file_names[mod2][idx_b]
+                ))
+                all_scores.append(score)
+    
+    # Compute statistics
+    if all_scores:
+        statistics.update({
+            "avg_score": float(np.mean(all_scores)),
+            "max_score": float(np.max(all_scores)),
+            "min_score": float(np.min(all_scores)),
+            "total_comparisons": len(all_scores)
+        })
+    
+    # Sort matches by score in descending order
+    matches.sort(key=lambda x: x.score, reverse=True)
+    
+    return SimilarityResponse(
+        matches=matches,
+        statistics=statistics,
+        modality_pairs=modality_pairs
+    )
 
 @app.get("/health")
 async def health_check(