Update app.py

app.py

@@ -78,6 +78,7 @@ def embedding_worker():
 
         embedding_response_queue.put(formatted_response)
         embedding_request_queue.task_done()
+        clear_cuda_cache()
 
 threading.Thread(target=embedding_worker, daemon=True).start()
 
@@ -100,6 +101,7 @@ def compute_embeddings(selected_task, input_text):
     embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
     embeddings = F.normalize(embeddings, p=2, dim=1)
     embeddings_list = embeddings.detach().cpu().numpy().tolist()
+    clear_cuda_cache()
     return embeddings_list
 
 @spaces.GPU
@@ -130,6 +132,7 @@ def compute_similarity(selected_task, sentence1, sentence2, extra_sentence1, ext
     free_memory(embeddings1, embeddings2, embeddings3, embeddings4)
 
     similarity_scores = {"Similarity 1-2": similarity1, "Similarity 1-3": similarity2, "Similarity 1-4": similarity3}
+    clear_cuda_cache()
     return similarity_scores
     
 @spaces.GPU
@@ -138,6 +141,7 @@ def compute_cosine_similarity(emb1, emb2):
     tensor2 = torch.tensor(emb2).to(device).half()
     similarity = F.cosine_similarity(tensor1, tensor2).item()
     free_memory(tensor1, tensor2)
+    clear_cuda_cache()
     return similarity
 
 
@@ -153,6 +157,7 @@ def compute_embeddings_batch(input_texts):
     outputs = model(**batch_dict)
     embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
     embeddings = F.normalize(embeddings, p=2, dim=1)
+    clear_cuda_cache()
     return embeddings.detach().cpu().numpy()
 
 def semantic_search(query_embedding, corpus_embeddings, top_k=5):
@@ -188,6 +193,7 @@ def generate_and_format_embeddings(selected_task, input_text):
     embedding_request_queue.put((selected_task, input_text))
     response = embedding_response_queue.get()
     embedding_response_queue.task_done()
+    clear_cuda_cache()
     return response