Update app.py

app.py

@@ -1,4 +1,4 @@
-import spaces
+# import spaces
 import torch
 import torch.nn.functional as F
 from torch import Tensor
@@ -9,7 +9,7 @@ import os
 title = """
 # 👋🏻Welcome to 🙋🏻‍♂️Tonic's 🐣e5-mistral🛌🏻Embeddings """
 description = """
-You can use this Space to test out the current model [intfloat/e5-mistral-7b-instruct](https://huggingface.co/intfloat/e5-mistral-7b-instruct). e5mistral has a larger context window, a different prompting/return mechanism and generally better results than other embedding models. 
+You can use this ZeroGPU Space to test out the current model [intfloat/e5-mistral-7b-instruct](https://huggingface.co/intfloat/e5-mistral-7b-instruct). 🐣e5-mistral🛌🏻 has a larger context🪟window, a different prompting/return🛠️mechanism and generally better results than other embedding models. use it via API to create embeddings or try out the sentence similarity to see how various optimization parameters affect performance.
 You can also use 🐣e5-mistral🛌🏻 by cloning this space. 🧬🔬🔍 Simply click here: <a style="display:inline-block" href="https://huggingface.co/spaces/Tonic/e5?duplicate=true"><img src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a></h3> 
 Join us : 🌟TeamTonic🌟 is always making cool demos! Join our active builder's🛠️community 👻  [![Join us on Discord](https://img.shields.io/discord/1109943800132010065?label=Discord&logo=discord&style=flat-square)](https://discord.gg/GWpVpekp) On 🤗Huggingface: [TeamTonic](https://huggingface.co/TeamTonic) & [MultiTransformer](https://huggingface.co/MultiTransformer) On 🌐Github: [Polytonic](https://github.com/tonic-ai) & contribute to 🌟 [Poly](https://github.com/tonic-ai/poly) 🤗Big thanks to Yuvi Sharma and all the folks at huggingface for the community grant 🤗
 """
@@ -44,52 +44,111 @@ 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()
 
-@spaces.GPU
-def compute_embeddings(selected_task, input_text, system_prompt):
-    max_length = 2042
-    task_description = tasks[selected_task]
-    processed_texts = [f'Instruct: {task_description}\nQuery: {input_text}']
+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
+    
+    def compute_embeddings(self, selected_task, input_text):
+        try:
+            task_description = tasks[selected_task]
+        except KeyError:
+            print(f"Selected task not found: {selected_task}")
+            return f"Error: Task '{selected_task}' not found. Please select a valid task."
+        max_length = 2042
+        processed_texts = [f'Instruct: {task_description}\nQuery: {input_text}']
+    
+        batch_dict = self.tokenizer(processed_texts, max_length=max_length - 1, return_attention_mask=False, padding=False, truncation=True)
+        batch_dict['input_ids'] = [input_ids + [self.tokenizer.eos_token_id] for input_ids in batch_dict['input_ids']]
+        batch_dict = self.tokenizer.pad(batch_dict, padding=True, return_attention_mask=True, return_tensors='pt')
+        batch_dict = {k: v.to(device) for k, v in batch_dict.items()}
+        outputs = self.model(**batch_dict)
+        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()
+        return embeddings_list
+
+    def compute_similarity(self, selected_task, sentence1, sentence2, extra_sentence1, extra_sentence2):
+        try:
+            task_description = tasks[selected_task]
+        except KeyError:
+            print(f"Selected task not found: {selected_task}")
+            return f"Error: Task '{selected_task}' not found. Please select a valid task."
+        # Compute embeddings for each sentence
+        embeddings1 = self.compute_embeddings(self.selected_task, sentence1)
+        embeddings2 = self.compute_embeddings(self.selected_task, sentence2)
+        embeddings3 = self.compute_embeddings(self.selected_task, extra_sentence1)
+        embeddings4 = self.compute_embeddings(self.selected_task, extra_sentence2)
+    
+        # Convert embeddings to tensors
+        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 = 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
 
-    batch_dict = tokenizer(processed_texts, max_length=max_length - 1, return_attention_mask=False, padding=False, truncation=True)
-    batch_dict['input_ids'] = [input_ids + [tokenizer.eos_token_id] for input_ids in batch_dict['input_ids']]
-    batch_dict = tokenizer.pad(batch_dict, padding=True, return_attention_mask=True, return_tensors='pt')
-    batch_dict = {k: v.to(device) for k, v in batch_dict.items()}
-    outputs = model(**batch_dict)
-    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()
-    return embeddings_list
     
 def app_interface():
+    embedding_model = EmbeddingModel()
     with gr.Blocks() as demo:
         gr.Markdown(title)
         gr.Markdown(description)
+        with gr.Row():
+            task_dropdown = gr.Dropdown(list(tasks.keys()), label="Select a Task", value=list(tasks.keys())[0])
+
+        with gr.Tab("Embedding Generation"):
+            input_text_box = gr.Textbox(label="📖Input Text")
+            compute_button = gr.Button("Try🐣🛌🏻e5")
+            output_display = gr.Textbox(label="🐣e5-mistral🛌🏻 Embeddings")
+            compute_button.click(
+                fn=embedding_model.compute_embeddings,
+                inputs=[task_dropdown, input_text_box],
+                outputs=output_display
+            )
+
+        with gr.Tab("Sentence Similarity"):
+            sentence1_box = gr.Textbox(label="'Focus Sentence' - The 'Subject'")
+            sentence2_box = gr.Textbox(label="'Input Sentence' - 1")
+            extra_sentence1_box = gr.Textbox(label="'Input Sentence' - 2")
+            extra_sentence2_box = gr.Textbox(label="'Input Sentence' - 3")
+            similarity_button = gr.Button("Compute Similarity")
+            similarity_output = gr.Label(label="🐣e5-mistral🛌🏻 Similarity Scores")
+            similarity_button.click(
+                fn=embedding_model.compute_similarity,
+                inputs=[task_dropdown, sentence1_box, sentence2_box],
+                outputs=similarity_output
+            )
 
-        task_dropdown = gr.Dropdown(list(tasks.keys()), label="Select a Task", value=list(tasks.keys())[0])
-     
-        input_text_box = gr.Textbox(label="📖Input Text")
-        system_prompt_box = gr.Textbox(label="🤖System Prompt (Optional)")
-        
-        compute_button = gr.Button("Try🐣🛌🏻e5")
-        
-        output_display = gr.Textbox(label="🐣e5-mistral🛌🏻")
-        
         with gr.Row():
             with gr.Column():
-                system_prompt_box
                 input_text_box
             with gr.Column():
                 compute_button
                 output_display
 
-        compute_button.click(
-            fn=compute_embeddings,
-            inputs=[task_dropdown, input_text_box],
-            outputs=output_display
-        )
-
-
     return demo
 
 # Run the Gradio app