import gradio as gr
from sentence_transformers import SentenceTransformer
import spaces
import torch
import ast

# Load the model
model = SentenceTransformer("Qwen/Qwen3-Embedding-0.6B")

@spaces.GPU
def compute(queries_input, documents_input):
    try:
        # Convert string input to Python lists
        queries = ast.literal_eval(queries_input)
        documents = ast.literal_eval(documents_input)

        # Validate input
        if not isinstance(queries, list) or not isinstance(documents, list):
            return "Inputs must be lists."

        # Generate embeddings
        query_embeddings = model.encode(queries, prompt_name="query")
        document_embeddings = model.encode(documents)

        # Compute similarity
        similarity_matrix = torch.tensor(model.similarity(query_embeddings, document_embeddings))
        return similarity_matrix.tolist()
    except Exception as e:
        return str(e)

demo = gr.Interface(
    fn=compute,
    inputs=[
        gr.Textbox(label="Queries (Python list, e.g. ['query1', 'query2'])"),
        gr.Textbox(label="Documents (Python list, e.g. ['doc1', 'doc2'])")
    ],
    outputs="json",
    title="embedding"
)
demo.launch()