Spaces:
Running
Running
File size: 3,012 Bytes
649d5f9 18c9386 649d5f9 18c9386 fd20dab ac5cb04 18c9386 649d5f9 18c9386 bea665a 57f1fb2 bea665a 57f1fb2 bea665a 18c9386 815dd17 |
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 |
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')
# dense_model = SentenceTransformer('distiluse-base-multilingual-cased-v2')
dense_model = SentenceTransformer('multi-qa-mpnet-base-cos-v1')
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])
unique_terms = sorted(set(tokens))
scores = bm25.get_scores(unique_terms)
# Assign scores for all unique terms
term_weights = {term: float(score) for term, score in zip(unique_terms, scores)}
indices = list(range(len(unique_terms)))
values = [term_weights.get(term, 0.0) for term in unique_terms]
return {"indices": indices, "values": values, "terms": unique_terms}
# 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(mcp_server=True)
|