File size: 1,193 Bytes
09ea753 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
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() |