Update app.py

app.py

@@ -44,12 +44,25 @@ def last_token_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tenso
         sequence_lengths = attention_mask.sum(dim=1) - 1
         batch_size = last_hidden_states.shape[0]
         return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]
+def clear_cuda_cache():
+    torch.cuda.empty_cache()
 
+def free_memory(*args):
+    for arg in args:
+        del arg
+        
 class EmbeddingModel:
     def __init__(self):
         self.tokenizer = AutoTokenizer.from_pretrained('intfloat/e5-mistral-7b-instruct')
         self.model = AutoModel.from_pretrained('intfloat/e5-mistral-7b-instruct', torch_dtype=torch.float16, device_map=device)
-
+    
+    def _compute_cosine_similarity(self, emb1, emb2):
+        tensor1 = torch.tensor(emb1).to(device).half()
+        tensor2 = torch.tensor(emb2).to(device).half()
+        similarity = F.cosine_similarity(tensor1, tensor2).item()
+        free_memory(tensor1, tensor2)
+        return similarity
+    
     @spaces.GPU
     def compute_embeddings(self, selected_task, input_text):
         try:
@@ -84,16 +97,20 @@ class EmbeddingModel:
         embeddings4 = self.compute_embeddings(self.selected_task, extra_sentence2)
     
         # Convert embeddings to tensors
-        embeddings_tensor1 = torch.tensor(embeddings1).to(device)
-        embeddings_tensor2 = torch.tensor(embeddings2).to(device)
-        embeddings_tensor3 = torch.tensor(embeddings3).to(device)
-        embeddings_tensor4 = torch.tensor(embeddings4).to(device)
+        embeddings_tensor1 = torch.tensor(embeddings1).to(device).half()
+        embeddings_tensor2 = torch.tensor(embeddings2).to(device).half()
+        embeddings_tensor3 = torch.tensor(embeddings3).to(device).half()
+        embeddings_tensor4 = torch.tensor(embeddings4).to(device).half()
     
         # Compute cosine similarity
-        similarity1 = F.cosine_similarity(embeddings_tensor1, embeddings_tensor2).item()
-        similarity2 = F.cosine_similarity(embeddings_tensor1, embeddings_tensor3).item()
-        similarity3 = F.cosine_similarity(embeddings_tensor1, embeddings_tensor4).item()
-        return similarity1, similarity2
+        similarity1 = self._compute_cosine_similarity(embeddings1, embeddings2)
+        similarity2 = self._compute_cosine_similarity(embeddings1, embeddings3)
+        similarity3 = self._compute_cosine_similarity(embeddings1, embeddings4)
+
+        # Free memory
+        free_memory(embeddings1, embeddings2, embeddings3, embeddings4)
+
+        return similarity1, similarity2, similarity3
 
     
 def app_interface():