normalize vectors

app.py

@@ -4,7 +4,7 @@ from src.encoder import FashionCLIPEncoder
 from src.models import TextRequest, ImageRequest, Response
 
 
-encoder = FashionCLIPEncoder()
+encoder = FashionCLIPEncoder(normalize=True)
 app = FastAPI()
 
 

src/encoder.py

@@ -4,12 +4,16 @@ from PIL.Image import Image
 import torch
 from transformers import AutoModel, AutoProcessor
 
+from .utils import normalize_vectors
+
 
 MODEL_NAME = "Marqo/marqo-fashionCLIP"
 
 
 class FashionCLIPEncoder:
-    def __init__(self):
+    def __init__(self, normalize: bool = False):
+        self.normalize = normalize
+
         self.device = torch.device("cpu")
 
         self.processor = AutoProcessor.from_pretrained(
@@ -31,24 +35,26 @@ class FashionCLIPEncoder:
             "return_tensors": "pt",
             "truncation": True,
         }
+
         inputs = self.processor(text=texts, **kwargs)
 
         with torch.no_grad():
             batch = {k: v.to(self.device) for k, v in inputs.items()}
-            return self._encode_text(batch)
+            vectors = self.model.get_text_features(**batch)
+                
+            return self._postprocess_vectors(vectors)
 
     def encode_images(self, images: List[Image]) -> List[List[float]]:
-        kwargs = {
-            "return_tensors": "pt",
-        }
-        inputs = self.processor(images=images, **kwargs)
+        inputs = self.processor(images=images, return_tensors="pt")
 
         with torch.no_grad():
             batch = {k: v.to(self.device) for k, v in inputs.items()}
-            return self._encode_images(batch)
+            vectors = self.model.get_image_features(**batch)
+
+            return self._postprocess_vectors(vectors)
 
-    def _encode_text(self, batch: Dict) -> List[List[float]]:
-        return self.model.get_text_features(**batch).detach().cpu().numpy().tolist()
+    def _postprocess_vectors(self, vectors: torch.Tensor) -> List[List[float]]:
+        if self.normalize:
+            vectors = normalize_vectors(vectors)
 
-    def _encode_images(self, batch: Dict) -> List[List[float]]:
-        return self.model.get_image_features(**batch).detach().cpu().numpy().tolist()
+        return vectors.detach().cpu().numpy().tolist()

src/utils.py

@@ -20,6 +20,15 @@ def download_image_as_pil(url: str, timeout: int = 10) -> Image.Image:
 
     except Exception as e:
         return
+    
+
+def normalize_vectors(vectors: torch.Tensor) -> torch.Tensor:
+    norms = torch.norm(vectors, p=2, dim=1, keepdim=True)
+    norms = torch.norm(vectors, p=2, dim=1, keepdim=True)
+    norms = torch.where(norms > 1e-8, norms, torch.ones_like(norms))
+    normalized_vectors = vectors / norms
+
+    return normalized_vectors
 
 
 def analyze_model_parameters(model: torch.nn.Module) -> Dict: