Spaces:
Running
Running
File size: 2,661 Bytes
649d5f9 18c9386 649d5f9 18c9386 649d5f9 18c9386 649d5f9 18c9386 |
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 |
import gradio as gr
from sentence_transformers import SentenceTransformer
from rank_bm25 import BM25Okapi
from transformers import AutoTokenizer, AutoModel
import torch
# 1. Dense embedding model (HF bi-encoder)
dense_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
def embed_dense(text: str):
if not text.strip():
return {"error": "Input text is empty."}
emb = dense_model.encode([text])[0]
return {"dense_embedding": emb.tolist()}
# 2. Sparse embedding model (BM25)
# Uses rank_bm25 to compute term weights
def embed_sparse(text: str):
if not text.strip():
return {"error": "Input text is empty."}
tokens = text.split()
bm25 = BM25Okapi([tokens])
scores = bm25.get_scores(tokens)
# Map each term to its BM25 weight
term_weights = {tok: float(score) for tok, score in zip(tokens, scores)}
return {"sparse_embedding": term_weights}
# 3. Late-interaction embedding model (ColBERT)
colbert_tokenizer = AutoTokenizer.from_pretrained('colbert-ir/colbertv2.0', use_fast=True)
colbert_model = AutoModel.from_pretrained('colbert-ir/colbertv2.0')
# Freeze model parameters for inference speed
for param in colbert_model.parameters():
param.requires_grad = False
def embed_colbert(text: str):
if not text.strip():
return {"error": "Input text is empty."}
inputs = colbert_tokenizer(text, return_tensors='pt', truncation=True, max_length=64)
with torch.no_grad():
outputs = colbert_model(**inputs)
# last_hidden_state: (1, seq_len, hidden_size)
embeddings = outputs.last_hidden_state.squeeze(0).tolist()
return {"colbert_embeddings": embeddings}
# Build Gradio interface with tabs for each model
with gr.Blocks(title="Text Embedding Playground") as demo:
gr.Markdown("# Text Embedding Playground\nChoose a model and input text to get embeddings.")
with gr.Tab("Dense (MiniLM-L6-v2)"):
txt1 = gr.Textbox(lines=3, label="Input Text")
out1 = gr.JSON(label="Embedding")
txt1.submit(embed_dense, txt1, out1)
gr.Button("Embed").click(embed_dense, txt1, out1)
with gr.Tab("Sparse (BM25)"):
txt2 = gr.Textbox(lines=3, label="Input Text")
out2 = gr.JSON(label="Term Weights")
txt2.submit(embed_sparse, txt2, out2)
gr.Button("Embed").click(embed_sparse, txt2, out2)
with gr.Tab("Late-Interaction (ColBERT)"):
txt3 = gr.Textbox(lines=3, label="Input Text")
out3 = gr.JSON(label="Embeddings per Token")
txt3.submit(embed_colbert, txt3, out3)
gr.Button("Embed").click(embed_colbert, txt3, out3)
if __name__ == "__main__":
demo.launch()
|