embed_text / app.py
bluuebunny's picture
Update app.py
e9ce7de verified
raw
history blame contribute delete
830 Bytes
from sentence_transformers import SentenceTransformer
import gradio as gr
import numpy as np
import torch
# Function to convert dense vector to binary vector
def dense_to_binary(dense_vector):
return np.packbits(np.where(dense_vector >= 0, 1, 0)).tobytes()
# Load the model
device="cuda" if torch.cuda.is_available() else "cpu"
model = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1", device=device)
def embed(text:str):
# Float embedding
float_vector = model.encode(text, convert_to_numpy=True)
# Convert to binary vector
binary_vector = dense_to_binary(float_vector)
# Return both vectors
return float_vector, binary_vector
# Gradio interface
interface = gr.Interface(
fn=embed,
inputs=["text"],
outputs=["json", "text"]
)
interface.launch(server_port=7860)