| import gradio as gr |
| import torch |
| from transformers import ( |
| AutoModel, |
| AutoTokenizer, |
| ) |
| import os |
| from threading import Thread |
| |
| import time |
|
|
| if torch.cuda.is_available(): |
| device = torch.device("cuda") |
| print(f"Using GPU: {torch.cuda.get_device_name(device)}") |
| else: |
| device = torch.device("cpu") |
| print("Using CPU") |
|
|
|
|
| def mean_pooling(model_output, attention_mask): |
| token_embeddings = model_output[0] |
| input_mask_expanded = ( |
| attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() |
| ) |
| return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp( |
| input_mask_expanded.sum(1), min=1e-9 |
| ) |
|
|
|
|
| def cls_pooling(model_output): |
| return model_output[0][:, 0] |
|
|
|
|
| |
| def get_embedding(text, use_mean_pooling, model_id): |
|
|
| tokenizer = AutoTokenizer.from_pretrained(model_id) |
| model = AutoModel.from_pretrained(model_id, torch_dtype=torch.float16) |
|
|
| model = model.to(device) |
| inputs = tokenizer( |
| text, return_tensors="pt", padding=True, truncation=True, max_length=512 |
| ) |
| inputs = {name: tensor.to(device) for name, tensor in inputs.items()} |
| with torch.no_grad(): |
| model_output = model(**inputs) |
| if use_mean_pooling: |
| return mean_pooling(model_output, inputs["attention_mask"]) |
| return cls_pooling(model_output) |
|
|
|
|
| def get_similarity(text1, text2, pooling_method, model_id): |
| use_mean_pooling = pooling_method == "Use Mean Pooling" |
| embedding1 = get_embedding(text1, use_mean_pooling, model_id) |
| embedding2 = get_embedding(text2, use_mean_pooling, model_id) |
| return torch.nn.functional.cosine_similarity(embedding1, embedding2).item() |
|
|
|
|
| gr.Interface( |
| get_similarity, |
| [ |
| gr.Textbox(lines=7, label="Text 1"), |
| gr.Textbox(lines=7, label="Text 2"), |
| gr.Dropdown( |
| choices=["Use Mean Pooling", "Use CLS"], |
| value="Use Mean Pooling", |
| label="Pooling Method", |
| info="Mean Pooling: Averages all token embeddings (better for semantic similarity)\nCLS Pooling: Uses only the [CLS] token embedding (faster, might miss context)", |
| ), |
| gr.Dropdown( |
| choices=[ |
| "nomic-ai/modernbert-embed-base", |
| "tasksource/ModernBERT-base-embed", |
| "tasksource/ModernBERT-base-nli", |
| "joe32140/ModernBERT-large-msmarco", |
| "answerdotai/ModernBERT-large", |
| "answerdotai/ModernBERT-base", |
| ], |
| value="answerdotai/ModernBERT-large", |
| label="Model", |
| info="Choose between the variants of ModernBERT \nMight take a few seconds to load the model", |
| ), |
| ], |
| gr.Textbox(label="Similarity"), |
| title="ModernBERT Similarity Demo", |
| description="Compute the similarity between two texts using ModernBERT. Choose between different pooling strategies for embedding generation.", |
| examples=[ |
| [ |
| "The quick brown fox jumps over the lazy dog", |
| "A swift brown fox leaps above a sleeping canine", |
| "Use Mean Pooling", |
| "answerdotai/ModernBERT-large", |
| ], |
| [ |
| "I love programming in Python", |
| "I hate coding with Python", |
| "Use Mean Pooling", |
| "joe32140/ModernBERT-large-msmarco", |
| ], |
| [ |
| "The weather is beautiful today", |
| "Machine learning models are improving rapidly", |
| "Use Mean Pooling", |
| "tasksource/ModernBERT-base-embed", |
| ], |
| [ |
| "def calculate_sum(a, b):\n return a + b", |
| "def add_numbers(x, y):\n result = x + y\n return result", |
| "Use Mean Pooling", |
| "tasksource/ModernBERT-base-nli", |
| ], |
| ], |
| ).launch() |
|
|
