File size: 3,825 Bytes
b8c24aa 3a82207 63b82b4 3592f5f 63b82b4 c8fdb3b 3a82207 df2517f 7dc3087 08c1bd3 64d8a64 63b82b4 64d8a64 63b82b4 64d8a64 63b82b4 08c1bd3 3592f5f 63b82b4 3592f5f 3a82207 3592f5f 3a82207 04c5c52 3592f5f 3a82207 3592f5f 63b82b4 3592f5f 63b82b4 3592f5f 46d2261 3592f5f 63b82b4 3592f5f 5772409 3592f5f 5772409 3592f5f 5772409 3592f5f 5772409 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 |
import gradio as gr
import torch
from transformers import (
AutoModel,
AutoTokenizer,
)
import os
from threading import Thread
# import spaces
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]
# @spaces.GPU
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()
|