Update app.py

app.py

@@ -11,16 +11,33 @@ title = """
 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 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 👻[![Let's build the future of AI together! 🚀🤖](https://discordapp.com/api/guilds/1109943800132010065/widget.png)](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 🤗
+Join us : 🌟TeamTonic🌟 is always making cool demos! Join our active builder's🛠️community 👻[![(https://discordapp.com/api/guilds/1109943800132010065/widget.png)](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 🤗
+
+You can use this space in **two ways !** either select an embeddings mode or 'None' to speak with the e5mistral LLM 🤗
 """
 
 os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:30'
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+tasks = {
+        'ArguAna': 'Given a claim, find documents that refute the claim',
+        'ClimateFEVER': 'Given a claim about climate change, retrieve documents that support or refute the claim',
+        'DBPedia': 'Given a query, retrieve relevant entity descriptions from DBPedia',
+        'FEVER': 'Given a claim, retrieve documents that support or refute the claim',
+        'FiQA2018': 'Given a financial question, retrieve user replies that best answer the question',
+        'HotpotQA': 'Given a multi-hop question, retrieve documents that can help answer the question',
+        'MSMARCO': 'Given a web search query, retrieve relevant passages that answer the query',
+        'NFCorpus': 'Given a question, retrieve relevant documents that best answer the question',
+        'NQ': 'Given a question, retrieve Wikipedia passages that answer the question',
+        'QuoraRetrieval': 'Given a question, retrieve questions that are semantically equivalent to the given question',
+        'SCIDOCS': 'Given a scientific paper title, retrieve paper abstracts that are cited by the given paper',
+        'SciFact': 'Given a scientific claim, retrieve documents that support or refute the claim',
+        'Touche2020': 'Given a question, retrieve detailed and persuasive arguments that answer the question',
+        'TRECCOVID': 'Given a query on COVID-19, retrieve documents that answer the query',
+}
+
 tokenizer = AutoTokenizer.from_pretrained('intfloat/e5-mistral-7b-instruct')
 model = AutoModel.from_pretrained('intfloat/e5-mistral-7b-instruct', torch_dtype=torch.float16, device_map=device)
-# model.half()
-# model.to(device)
 
 def last_token_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
     left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
@@ -31,17 +48,15 @@ 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]
 
-def get_detailed_instruct(task_description: str, query: str) -> str:
-    return f'Instruct: {task_description}\nQuery: {query}'
-
-    
 @spaces.GPU
-def compute_embeddings(*input_texts):
-
+def compute_embeddings(selected_task, *input_texts):
     max_length = 2042
-    task = 'Given a web search query, retrieve relevant passages that answer the query'
-
-    processed_texts = [get_detailed_instruct(task, text) for text in input_texts]
+    if selected_task:
+        task = tasks[selected_task]
+        processed_texts = [f'Instruct: {task}\nQuery: {text}' for text in input_texts]
+    else:
+        processed_texts = [f'Instruct: {system_prompt}\nQuerry: {text}' for text in input_texts]
+    task = tasks[selected_task]
     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')
@@ -56,12 +71,14 @@ def app_interface():
     with gr.Blocks() as demo:
         gr.Markdown(title)
         gr.Markdown(description)
+
+        task_dropdown = gr.Dropdown(list(tasks.keys()) + ["None"], label="Select a Task (Optional)", value="None")
         
-        input_text_boxes = [gr.Textbox(label=f"Input Text {i+1}") for i in range(4)]
+        input_text_boxes = gr.Textbox(label=f"Input Text")
         
-        compute_button = gr.Button("Compute Embeddings")
+        compute_button = gr.Button("Try🐣🛌🏻e5")
         
-        output_display = gr.Dataframe(headers=["Embedding Value"], datatype=["number"])
+        output_display = gr.Textbox(label="🐣e5-mistral🛌🏻")
         
         with gr.Row():
             with gr.Column():
@@ -73,7 +90,7 @@ def app_interface():
 
         compute_button.click(
             fn=compute_embeddings,
-            inputs=input_text_boxes,
+            inputs=[task_dropdown] + input_text_boxes,
             outputs=output_display
         )