Update app.py

app.py

@@ -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 🤗
 """
@@ -45,36 +45,75 @@ def last_token_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tenso
         batch_size = last_hidden_states.shape[0]
         return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]
 
-@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}']
+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)
+
+    @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}']
+    
+        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
+
+    @spaces.GPU
+    def compute_similarity(self, sentence1, sentence2, extra_sentence1, extra_sentence2):
+        # Tokenize and encode sentences
+        sentences = [sentence1, sentence2, extra_sentence1, extra_sentence2]
+        encoded_input = self.tokenizer(sentences, padding=True, truncation=True, return_tensors='pt').to(device)
+        with torch.no_grad():
+            model_output = self.model(**encoded_input)
+
+        # Compute embeddings
+        embeddings = last_token_pool(model_output.last_hidden_state, encoded_input['attention_mask'])
+        embeddings = F.normalize(embeddings, p=2, dim=1)
+
+        # Compute cosine similarity
+        similarity1 = F.cosine_similarity(embeddings[0].unsqueeze(0), embeddings[1].unsqueeze(0)).item()
+        similarity2 = F.cosine_similarity(embeddings[2].unsqueeze(0), embeddings[3].unsqueeze(0)).item()
+        return similarity1, similarity2
 
-    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():
     with gr.Blocks() as demo:
         gr.Markdown(title)
         gr.Markdown(description)
 
-        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.Tab("Embedding Generation"):
+            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🛌🏻 Embeddings")
+            compute_button.click(
+                fn=embedding_model.compute_embeddings,
+                inputs=[task_dropdown, input_text_box, system_prompt_box],
+                outputs=output_display
+            )
+
+        with gr.Tab("Sentence Similarity"):
+            sentence1_box = gr.Textbox(label="Sentence 1")
+            sentence2_box = gr.Textbox(label="Sentence 2")
+            extra_sentence1_box = gr.Textbox(label="Extra Sentence 1")
+            extra_sentence2_box = gr.Textbox(label="Extra Sentence 2")
+            similarity_button = gr.Button("Compute Similarity")
+            similarity_output = gr.Label(label="🐣e5-mistral🛌🏻 Similarity Scores")
+            similarity_button.click(
+                fn=embedding_model.compute_similarity,
+                inputs=[sentence1_box, sentence2_box, extra_sentence1_box, extra_sentence2_box],
+                outputs=similarity_output
+            )
+
         with gr.Row():
             with gr.Column():
                 system_prompt_box
@@ -83,13 +122,6 @@ def app_interface():
                 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