import gradio as gr
from transformers import AutoTokenizer, AutoModel
import torch
from sentence_transformers import SentenceTransformer


# Load your model with trust_remote_code=True
model = SentenceTransformer("truong1301/bi-encode-HG-DOCS", trust_remote_code=True)

# Define a function to get embeddings
def get_embedding(text):
    with torch.no_grad():
        embeddings = model.encode(text)
    return embeddings

# Create a Gradio interface
iface = gr.Interface(
    fn=get_embedding,
    inputs=gr.Textbox(lines=2, placeholder="Enter text here..."),
    outputs="json",
    title="Embedding Generator",
    description="Get embeddings using truong1301/bi-encode-HG-DOCS"
)

# Launch the Gradio app
iface.launch()