Update app.py

app.py

@@ -1,41 +1,66 @@
 import gradio as gr
 from sentence_transformers import SentenceTransformer
+import torch
 
-# Load a fast, high-quality embedding model (runs great on CPU or GPU)
-model = SentenceTransformer("all-MiniLM-L6-v2")   # 22M params, ~80 MB
-# Alternatives: "BAAI/bge-small-en-v1.5", "nomic-ai/nomic-embed-text-v1", etc.
+# Load Qwen3-Embedding-4B (best performance on GPU)
+model = SentenceTransformer(
+    "Qwen/Qwen3-Embedding-4B",
+    model_kwargs={
+        "attn_implementation": "flash_attention_2",  # recommended for speed + lower VRAM
+        "device_map": "auto",
+        "torch_dtype": torch.float16,                # halves memory usage
+    },
+    tokenizer_kwargs={"padding_side": "left"},
+)
 
-def generate_embeddings(texts: str):
-    """Accepts one text per line (batch) or a single sentence."""
-    # Split by lines for easy batching
+def generate_embeddings(texts: str, use_query_mode: bool):
+    """One text per line. Returns 2560-dim embeddings."""
     lines = [line.strip() for line in texts.split("\n") if line.strip()]
     if not lines:
-        return {"embeddings": [], "texts": []}
-    
-    # Generate embeddings (automatically batched)
-    embeddings = model.encode(lines, convert_to_numpy=True).tolist()
-    
+        return {"embeddings": [], "texts": [], "dimension": 2560}
+
+    # Use "query" prompt for better retrieval performance (adds instruction internally)
+    prompt_name = "query" if use_query_mode else None
+
+    embeddings = model.encode(
+        lines,
+        prompt_name=prompt_name,
+        convert_to_numpy=True,
+        normalize_embeddings=True,   # cosine-ready vectors
+    ).tolist()
+
     return {
-        "embeddings": embeddings,   # list of lists of floats
-        "texts": lines
+        "embeddings": embeddings,    # list of lists (2560 floats each)
+        "texts": lines,
+        "dimension": len(embeddings[0]) if embeddings else 2560,
+        "mode": "query (with instruction)" if use_query_mode else "document (raw)",
     }
 
-# Gradio interface (provides both a nice UI for testing + the full API)
+
+# Gradio UI + automatic REST API
 demo = gr.Interface(
     fn=generate_embeddings,
-    inputs=gr.Textbox(
-        lines=8,
-        placeholder="Enter one text per line for batch embedding...\n\nThe quick brown fox jumps over the lazy dog.",
-        label="Input Texts"
-    ),
-    outputs=gr.JSON(label="Embeddings"),
-    title="🔥 Free Embedding API",
-    description="Powered by sentence-transformers on Hugging Face Spaces. Returns 384-dimensional embeddings (MiniLM).",
+    inputs=[
+        gr.Textbox(
+            lines=10,
+            placeholder="Paste one text per line...\n\nThe capital of France is Paris.\nExplain quantum entanglement.",
+            label="Input Texts (one per line)",
+        ),
+        gr.Checkbox(
+            label="Use Query Mode (recommended for search/retrieval)",
+            value=True,
+            info="Adds task-specific instruction automatically for much better results",
+        ),
+    ],
+    outputs=gr.JSON(label="Embedding Response"),
+    title="🚀 Qwen3-Embedding-4B API",
+    description="""State-of-the-art 4B embedding model (2560-dim, 32k context, 100+ languages).
+    Powered by Qwen/Qwen3-Embedding-4B on Hugging Face Spaces.""",
     examples=[
-        ["The quick brown fox jumps over the lazy dog."],
-        ["Hello world!\nThis is a second sentence."]
+        ["What is the capital of China?\nExplain gravity", True],
+        ["The capital of China is Beijing.\nGravity is a force...", False],
     ],
-    allow_flagging="never"
+    allow_flagging="never",
 )
 
 demo.launch()