Update app.py

app.py

@@ -2,9 +2,9 @@ import os
 import torch
 import numpy as np
 from transformers import AutoTokenizer, AutoModel
-from flask import Flask, request, jsonify
+import gradio as gr
 import logging
-import spaces  # HuggingFace Spaces GPU decorator
+import spaces
 
 # Try to import flash attention (optional)
 try:
@@ -19,8 +19,6 @@ except ImportError:
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-app = Flask(__name__)
-
 # Qwen3-Embedding-4B model for retrieval
 MODEL_NAME = "Qwen/Qwen3-Embedding-4B"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@@ -31,7 +29,7 @@ tokenizer = None
 model = None
 
 def initialize_model():
-    """Initialize model (runs on CPU in main process)"""
+    """Initialize model"""
     global tokenizer, model
     
     if tokenizer is None:
@@ -44,7 +42,6 @@ def initialize_model():
     if model is None:
         logger.info(f"Loading {MODEL_NAME} on {DEVICE}")
         
-        # Configure model loading with optional flash attention
         model_kwargs = {
             "trust_remote_code": True,
             "torch_dtype": torch.float16 if DEVICE == "cuda" else torch.float32
@@ -59,12 +56,15 @@ def initialize_model():
         model.eval()
         logger.info("✅ Model loaded successfully")
 
-# CRITICAL: This must be a TOP-LEVEL function with @spaces.GPU decorator
 @spaces.GPU
-def encode_texts_gpu(texts, batch_size=16):
+def encode_texts_gpu(texts_str, batch_size=16):
     """
     Encode texts to embeddings using Qwen3-Embedding-4B
-    This function MUST be at module level for ZeroGPU detection
+    Args:
+        texts_str: Either a single text string or multiple texts separated by '|||'
+        batch_size: Batch size for encoding
+    Returns:
+        JSON string with embeddings
     """
     global tokenizer, model
     
@@ -72,8 +72,11 @@ def encode_texts_gpu(texts, batch_size=16):
     if model is None or tokenizer is None:
         initialize_model()
     
-    if isinstance(texts, str):
-        texts = [texts]
+    # Parse input - support both single text and multiple texts
+    if '|||' in texts_str:
+        texts = [t.strip() for t in texts_str.split('|||')]
+    else:
+        texts = [texts_str]
     
     embeddings = []
     
@@ -93,7 +96,6 @@ def encode_texts_gpu(texts, batch_size=16):
         
         with torch.no_grad():
             outputs = model(**inputs)
-            # Use EOS token embedding for Qwen3
             eos_token_id = tokenizer.eos_token_id
             sequence_lengths = (inputs['input_ids'] == eos_token_id).long().argmax(-1) - 1
             
@@ -103,82 +105,73 @@ def encode_texts_gpu(texts, batch_size=16):
                 batch_embeddings.append(embedding)
             
             batch_embeddings = np.array(batch_embeddings)
-            
-            # Normalize embeddings
             batch_embeddings = batch_embeddings / np.linalg.norm(batch_embeddings, axis=1, keepdims=True)
             
             embeddings.extend(batch_embeddings)
     
-    return embeddings
-
-@app.route("/", methods=["GET"])
-def health_check():
-    return jsonify({
-        "status": "healthy",
+    # Format output
+    import json
+    result = {
+        "embeddings": [emb.tolist() for emb in embeddings],
         "model": MODEL_NAME,
-        "device": DEVICE,
-        "embedding_dim": EMBEDDING_DIM,
-        "max_context": 32768
-    })
-
-@app.route("/embed", methods=["POST"])
-def embed_texts():
-    """Embed texts and return embeddings"""
-    try:
-        data = request.get_json()
-        
-        if not data or "texts" not in data:
-            return jsonify({"error": "Missing 'texts' field"}), 400
-        
-        texts = data["texts"]
-        if not isinstance(texts, list):
-            texts = [texts]
-        
-        logger.info(f"Embedding {len(texts)} texts")
-        
-        # Call the GPU-decorated function
-        embeddings = encode_texts_gpu(texts)
-        
-        return jsonify({
-            "embeddings": [embedding.tolist() for embedding in embeddings],
-            "model": MODEL_NAME,
-            "dimension": len(embeddings[0]) if embeddings else 0,
-            "count": len(embeddings)
-        })
-        
-    except Exception as e:
-        logger.error(f"Embedding error: {str(e)}")
-        return jsonify({"error": str(e)}), 500
+        "dimension": len(embeddings[0]) if embeddings else 0,
+        "count": len(embeddings)
+    }
+    
+    return json.dumps(result, indent=2)
 
-@app.route("/embed_single", methods=["POST"])
-def embed_single():
-    """Embed single text (convenience endpoint)"""
-    try:
-        data = request.get_json()
-        
-        if not data or "text" not in data:
-            return jsonify({"error": "Missing 'text' field"}), 400
-        
-        text = data["text"]
-        logger.info(f"Embedding single text: {text[:100]}...")
-        
-        # Call the GPU-decorated function
-        embeddings = encode_texts_gpu([text])
-        
-        return jsonify({
-            "embedding": embeddings[0].tolist(),
-            "model": MODEL_NAME,
-            "dimension": len(embeddings[0]),
-            "text_length": len(text)
-        })
-        
-    except Exception as e:
-        logger.error(f"Single embedding error: {str(e)}")
-        return jsonify({"error": str(e)}), 500
+# Create Gradio interface
+with gr.Blocks(title="Qwen3-Embedding-4B API") as demo:
+    gr.Markdown("""
+    # Qwen3-Embedding-4B Embedding Service
+    
+    This service generates embeddings using Qwen3-Embedding-4B (2560 dimensions).
+    
+    **Usage:**
+    - Single text: Enter your text directly
+    - Multiple texts: Separate texts with `|||` (e.g., `text1|||text2|||text3`)
+    """)
+    
+    with gr.Row():
+        with gr.Column():
+            text_input = gr.Textbox(
+                label="Text Input",
+                placeholder="Enter text or multiple texts separated by '|||'",
+                lines=5
+            )
+            batch_size_input = gr.Slider(
+                minimum=1,
+                maximum=64,
+                value=16,
+                step=1,
+                label="Batch Size"
+            )
+            submit_btn = gr.Button("Generate Embeddings", variant="primary")
+        
+        with gr.Column():
+            output = gr.JSON(label="Embeddings Output")
+    
+    submit_btn.click(
+        fn=encode_texts_gpu,
+        inputs=[text_input, batch_size_input],
+        outputs=output
+    )
+    
+    gr.Markdown("""
+    ### API Usage
+    You can also call this Space via API:
+    ```
+    from gradio_client import Client
+    
+    client = Client("YOUR_USERNAME/YOUR_SPACE_NAME")
+    result = client.predict(
+        texts_str="Your text here",
+        batch_size=16,
+        api_name="/predict"
+    )
+    print(result)
+    ```
+    """)
 
 if __name__ == "__main__":
-    logger.info("🚀 Starting embedding service...")
-    logger.info("⚡ Model will load on first GPU request (ZeroGPU lazy loading)")
-    
-    port = int(os.environ.get("PORT", 7860))
-    app.run(host="0.0.0.0", port=port)
+    demo.launch(server_name="0.0.0.0", server_port=7860)