import gradio as gr
from sentence_transformers import SentenceTransformer

# Load the Jina code-embedding model
model = SentenceTransformer(
    "jinaai/jina-embeddings-v2-base-code",
    trust_remote_code=True
)

def embed_code(code_snippet):
    """
    Given a string containing code, return its embedding vector as a list.
    """
    # The model.encode() method returns a numpy array; convert to list for display
    embedding = model.encode(code_snippet)
    return embedding.tolist()

# Build a simple Gradio interface
iface = gr.Interface(
    fn=embed_code,
    inputs=gr.Textbox(
        lines=10,
        placeholder="Paste your code snippet here..."
    ),
    outputs="textbox",
    title="Jina Code Embeddings",
    description="Embeds a code snippet into a vector using jina-embeddings-v2-base-code."
)

if __name__ == "__main__":
    iface.launch()