import gradio as gr
import spaces
from sentence_transformers import SentenceTransformer
from sentence_transformers.util import cos_sim
from sentence_transformers.quantization import quantize_embeddings

print("Loading embedding model");
dimensions = 768
model = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1", truncate_dim=dimensions)

@spaces.GPU
def embed(text):
        
    query_embedding = model.encode(text, prompt_name="query")
    return query_embedding.tolist();



with gr.Blocks() as demo:
    txtEmbed  = gr.Text(label="Text to embed")
    btnEmbed = gr.Button("embed");
    
    search = gr.Text(label="Script to search")
    
    results = gr.Text(label="results");
    
    btnEmbed.click(embed, [txtEmbed], [results])
    
    



if __name__ == "__main__":
    demo.launch(
        share=False,
        debug=False,
        server_port=7860,
        server_name="0.0.0.0",
        allowed_paths=[]
    )