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)