Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from sentence_transformers import SentenceTransformer, util | |
| import torch | |
| # Load the model (Qwen3-Embedding-0.6B) | |
| # trust_remote_code is required for some Qwen architectures | |
| model = SentenceTransformer("Qwen/Qwen3-Embedding-0.6B", trust_remote_code=True) | |
| def embed_text(text, is_query=True): | |
| """ | |
| API 1: Text Embedding | |
| Qwen3 benefits from a 'query' prompt for retrieval tasks. | |
| """ | |
| prompt_name = "query" if is_query else None | |
| embedding = model.encode(text, prompt_name=prompt_name, convert_to_tensor=True) | |
| return embedding.tolist() | |
| def calculate_similarity(text_a, text_b, is_query=True): | |
| """ | |
| API 2: Embedding Similarity | |
| Returns a float between 0 and 1 (clamped) representing the similarity. | |
| """ | |
| prompt_name = "query" if is_query else None | |
| # Encode both texts | |
| emb_a = model.encode(text_a, prompt_name=prompt_name, convert_to_tensor=True) | |
| emb_b = model.encode(text_b, prompt_name=prompt_name, convert_to_tensor=True) | |
| # Compute Cosine Similarity | |
| similarity = util.cos_sim(emb_a, emb_b).item() | |
| # Clamp to [0, 1] for "percentage" logic | |
| score = max(0, min(1, similarity)) | |
| percentage = f"{score * 100:.2f}%" | |
| return score, percentage | |
| # Building the Gradio UI | |
| with gr.Blocks(title="Qwen3 Embedding API") as demo: | |
| gr.Markdown("# Qwen3-Embedding-0.6B API & UI") | |
| gr.Markdown("This space provides high-quality text embeddings and similarity scores.") | |
| with gr.Tab("Text Embedding"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| input_text = gr.Textbox(label="Input Text", placeholder="Enter text to embed...") | |
| is_query_toggle = gr.Checkbox(label="Is this a search query?", value=True, info="Uses Qwen's specific query prompt for better retrieval.") | |
| btn_embed = gr.Button("Generate Embedding", variant="primary") | |
| with gr.Column(): | |
| output_vec = gr.JSON(label="Embedding Vector (Truncated in UI)") | |
| btn_embed.click(fn=embed_text, inputs=[input_text, is_query_toggle], outputs=output_vec, api_name="embed") | |
| with gr.Tab("Similarity Score"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| text_a = gr.Textbox(label="Text A", placeholder="First sentence...") | |
| text_b = gr.Textbox(label="Text B", placeholder="Second sentence...") | |
| is_query_sim = gr.Checkbox(label="Use query prompts?", value=True) | |
| btn_sim = gr.Button("Compare Texts", variant="primary") | |
| with gr.Column(): | |
| sim_float = gr.Number(label="Similarity Score (0-1)") | |
| sim_percent = gr.Label(label="Match Percentage") | |
| # API returns both the float and the label string | |
| btn_sim.click(fn=calculate_similarity, inputs=[text_a, text_b, is_query_sim], outputs=[sim_float, sim_percent], api_name="similarity") | |
| gr.Markdown(""" | |
| ### How to use the API | |
| You can call these endpoints programmatically using the Gradio Python Client: | |
| ```python | |
| from gradion_client import Client | |
| client = Client("your-username/your-space-name") | |
| # API 1: Embedding | |
| result = client.predict("Hello world", True, api_name="/embed") | |
| # API 2: Similarity | |
| score, percent = client.predict("Text A", "Text B", True, api_name="/similarity") | |
| ``` | |
| """) | |
| if __name__ == "__main__": | |
| demo.launch() |