|
|
import gradio as gr |
|
|
import torch |
|
|
from transformers import AutoTokenizer, AutoModel |
|
|
|
|
|
MODEL_NAME = "BAAI/bge-multilingual-gemma2" |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) |
|
|
model = AutoModel.from_pretrained(MODEL_NAME) |
|
|
|
|
|
def embed(text): |
|
|
inputs = tokenizer( |
|
|
text, |
|
|
return_tensors="pt", |
|
|
padding=True, |
|
|
truncation=True |
|
|
) |
|
|
with torch.no_grad(): |
|
|
outputs = model(**inputs) |
|
|
embeddings = outputs.last_hidden_state[:, 0] |
|
|
return embeddings[0].tolist() |
|
|
|
|
|
demo = gr.Interface( |
|
|
fn=embed, |
|
|
inputs=gr.Textbox(lines=4, placeholder="Enter text in any language"), |
|
|
outputs="json", |
|
|
title="BAAI/bge-multilingual-gemma2 Embedding Space", |
|
|
description="Multilingual embedding model for semantic search & RAG" |
|
|
) |
|
|
|
|
|
demo.launch() |
|
|
|