|
|
import spaces |
|
|
import torch |
|
|
import torch.nn.functional as F |
|
|
from torch import Tensor |
|
|
from transformers import AutoTokenizer, AutoModel |
|
|
import gradio as gr |
|
|
import os |
|
|
|
|
|
title = """ |
|
|
# 👋🏻Welcome to 🙋🏻♂️Tonic's 🐣e5-mistral🛌🏻Embeddings """ |
|
|
description = """ |
|
|
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 👻 [](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 🤗 |
|
|
""" |
|
|
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) |
|
|
|
|
|
def last_token_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor: |
|
|
left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0]) |
|
|
if left_padding: |
|
|
return last_hidden_states[:, -1] |
|
|
else: |
|
|
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 |
|
|
|
|
|
@spaces.GPU |
|
|
def compute_embeddings(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 = 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(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." |
|
|
|
|
|
embeddings1 = compute_embeddings(selected_task, sentence1) |
|
|
embeddings2 = compute_embeddings(selected_task, sentence2) |
|
|
embeddings3 = compute_embeddings(selected_task, extra_sentence1) |
|
|
embeddings4 = compute_embeddings(selected_task, extra_sentence2) |
|
|
|
|
|
|
|
|
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() |
|
|
|
|
|
|
|
|
similarity1 = compute_cosine_similarity(embeddings1, embeddings2) |
|
|
similarity2 = compute_cosine_similarity(embeddings1, embeddings3) |
|
|
similarity3 = compute_cosine_similarity(embeddings1, embeddings4) |
|
|
|
|
|
|
|
|
free_memory(embeddings1, embeddings2, embeddings3, embeddings4) |
|
|
|
|
|
similarity_scores = {"Similarity 1-2": similarity1, "Similarity 1-3": similarity2, "Similarity 1-4": similarity3} |
|
|
|
|
|
@spaces.GPU |
|
|
def compute_cosine_similarity(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 app_interface(): |
|
|
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=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=compute_similarity, |
|
|
inputs=[task_dropdown, sentence1_box, sentence2_box, extra_sentence1_box, extra_sentence2_box], |
|
|
outputs=similarity_output |
|
|
) |
|
|
|
|
|
with gr.Row(): |
|
|
with gr.Column(): |
|
|
input_text_box |
|
|
with gr.Column(): |
|
|
compute_button |
|
|
output_display |
|
|
|
|
|
return demo |
|
|
|
|
|
|
|
|
app_interface().launch() |
