# main.py import spaces import torch import torch.nn.functional as F from torch.nn import DataParallel from torch import Tensor from transformers import AutoTokenizer, AutoModel import threading import queue import gradio as gr import os import json import numpy as np title = """ # 👋🏻Welcome to 🙋🏻‍♂️Tonic's 📽️Nvidia 🛌🏻Embed V-1 !""" description = """ You can use this Space to test out the current model [nvidia/NV-Embed-v1](https://huggingface.co/nvidia/NV-Embed-v1). 🐣a generalist embedding model that ranks No. 1 on the Massive Text Embedding Benchmark (MTEB benchmark)(as of May 24, 2024), with 56 tasks, encompassing retrieval, reranking, classification, clustering, and semantic textual similarity tasks. You can also use 📽️Nvidia 🛌🏻Embed V-1 by cloning this space. 🧬🔬🔍 Simply click here: Duplicate Space 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: [Tonic-AI](https://github.com/tonic-ai) & contribute to 🌟 [MultiTonic](https://github.com/MultiTonic) 🤗Big thanks to Yuvi Sharma and all the folks at huggingface for the community grant 🤗 """ 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', 'Natural Language Inference' : 'Retrieve semantically similar text', 'Natural Language Inference' : 'Given a premise, retrieve a hypothesis that is entailed by the premise 20k', 'PAQ, MSMARCO' : 'Given a web search query, retrieve relevant passages that answer the query', 'PAQ, MSMARCO' : 'Given a question, retrieve passages that answer the question', 'SQUAD' : 'Given a question, retrieve Wikipedia passages that answer the question' , 'StackExchange' : 'Given a question paragraph at StackExchange, retrieve a question duplicated paragraph', 'Natural Question' : 'Given a question, retrieve Wikipedia passages that answer the question', 'BioASQ' : 'Given a question, retrieve detailed question descriptions that are duplicates to the given question', 'STS12, STS22, STSBenchmark' : 'Retrieve semantically similar text.', 'AmazonCounterfactualClassification' : 'Classify a given Amazon customer review text as either counterfactual or not-counterfactual' , 'AmazonReviewsClassification' : 'Classify the given Amazon review into its appropriate rating category' , 'Banking77Classification' : 'Given a online banking query, find the corresponding intents', 'EmotionClassification' : 'Classify the emotion expressed in the given Twitter message into one of the six emotions:anger, fear, joy, love, sadness, and surprise', 'ImdbClassification': 'Classify the sentiment expressed in the given movie review text from the IMDB dataset', 'MTOPIntentClassification' : 'Classify the intent of the given utterance in task-oriented conversation', 'ToxicConversationsClassification' : 'Classify the given comments as either toxic or not toxic', 'TweetSentimentExtractionClassification' : 'Classify the sentiment of a given tweet as either positive, negative, or neutral', 'ArxivClusteringP2P' : 'Identify the main and secondary category of Arxiv papers based on the titles and abstracts', 'ArxivClusteringS2S' : 'Identify the main and secondary category of Arxiv papers based on the titles', 'BiorxivClusteringP2P' : 'Identify the main category of Biorxiv papers based on the titles and abstracts' , 'BiorxivClusteringS2S' : 'Identify the main category of Biorxiv papers based on the titles', 'MedrxivClusteringP2P' : 'Identify the main category of Medrxiv papers based on the titles and abstracts', 'MedrxivClusteringS2S' : 'Identify the main category of Medrxiv papers based on the titles', 'TwentyNewsgroupsClustering' : 'Identify the topic or theme of the given news articles' } os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:30' os.environ['CUDA_LAUNCH_BLOCKING'] = '1' os.environ['CUDA_CACHE_DISABLE'] = '1' device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Load the tokenizer and model tokenizer = AutoTokenizer.from_pretrained('nvidia/NV-Embed-v1', trust_remote_code=True) model = AutoModel.from_pretrained('nvidia/NV-Embed-v1', trust_remote_code=True).to(device) # If multiple GPUs are available, enable DataParallel if torch.cuda.device_count() > 1: print("Using", torch.cuda.device_count(), "GPUs") model = DataParallel(model) # Embedding requests and response queues embedding_request_queue = queue.Queue() embedding_response_queue = queue.Queue() def clear_cuda_cache(): torch.cuda.empty_cache() 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 format_response(embeddings): return { "data": [ { "embedding": embeddings, "index": 0, "object": "embedding" } ], "model": "e5-mistral", "object": "list", "usage": { "prompt_tokens": 17, "total_tokens": 17 } } def embedding_worker(): while True: # Wait for an item in the queue item = embedding_request_queue.get() if item is None: break selected_task, input_text = item embeddings = compute_embeddings(selected_task, input_text) formatted_response = format_response(embeddings) embedding_response_queue.put(formatted_response) embedding_request_queue.task_done() clear_cuda_cache() 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." query_prefix = f"Instruct: {task_description}\nQuery: " queries = [input_text] # Get the embeddings query_embeddings = model.encode(queries, instruction=query_prefix, max_length=4096) # Normalize embeddings query_embeddings = F.normalize(query_embeddings, p=2, dim=1) embeddings_list = query_embeddings.detach().cpu().numpy().tolist() clear_cuda_cache() 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." # Compute embeddings for each sentence 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) similarity1 = compute_cosine_similarity(embeddings1, embeddings2) similarity2 = compute_cosine_similarity(embeddings1, embeddings3) similarity3 = compute_cosine_similarity(embeddings1, embeddings4) similarity_scores = {"Similarity 1-2": similarity1, "Similarity 1-3": similarity2, "Similarity 1-4": similarity3} clear_cuda_cache() return similarity_scores @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() clear_cuda_cache() 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("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.Textbox(label="📽️Nvidia 🛌🏻Embed Similarity Scores") similarity_button.click( fn=compute_similarity, inputs=[task_dropdown, sentence1_box, sentence2_box, extra_sentence1_box, extra_sentence2_box], outputs=similarity_output ) return demo embedding_worker_thread = threading.Thread(target=embedding_worker, daemon=True) embedding_worker_thread.start() app_interface().queue() app_interface().launch(share=True)