Upload serve_embedding.py with huggingface_hub

serve_embedding.py

@@ -0,0 +1,209 @@
+"""
+DeepX Embedding Server — Deploy on RTX 3060 12GB.
+
+Provides REST API for text embedding.
+Loads INT8 quantized model for efficient inference.
+
+Usage:
+  python scripts/serve_embedding.py \
+    --checkpoint checkpoints/deploy/deepx_int8.pt \
+    --port 8080
+
+API:
+  POST /embed
+  Body: {"texts": ["text1", "text2", ...], "normalize": true}
+  Response: {"embeddings": [[...], [...]], "dim": 1536}
+
+  POST /similarity
+  Body: {"query": "...", "documents": ["doc1", "doc2", ...]}
+  Response: {"scores": [0.85, 0.72, ...]}
+
+  GET /health
+  Response: {"status": "ok", "model": "deepx-v0.7", "device": "cuda"}
+"""
+import sys, os
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+
+import torch
+import torch.nn.functional as F
+import argparse
+import logging
+import time
+from typing import List
+from flask import Flask, request, jsonify
+from transformers import AutoTokenizer
+from config import DeepXConfig
+from modeling.pipeline import DeepXPipeline
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
+logger = logging.getLogger(__name__)
+
+app = Flask(__name__)
+
+# Globals
+pipeline = None
+tokenizer = None
+device = None
+MAX_LEN = 2048
+MAX_BATCH = 32
+
+
+def dequantize_int8(state_dict):
+    """Dequantize int8 back to float16."""
+    dequantized = {}
+    scale_keys = {k for k in state_dict if k.endswith("._scale")}
+    for key, tensor in state_dict.items():
+        if key in scale_keys:
+            continue
+        scale_key = key + "._scale"
+        if tensor.dtype == torch.int8 and scale_key in state_dict:
+            scale = state_dict[scale_key].float()
+            dequantized[key] = (tensor.float() * scale).half()
+        else:
+            dequantized[key] = tensor
+    return dequantized
+
+
+def load_model(checkpoint_path, tokenizer_path, device_str="cuda"):
+    global pipeline, tokenizer, device
+    
+    device = torch.device(device_str if torch.cuda.is_available() else "cpu")
+    logger.info(f"Device: {device}")
+    
+    # Tokenizer
+    logger.info(f"Loading tokenizer from {tokenizer_path}...")
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    
+    # Model
+    logger.info(f"Loading model from {checkpoint_path}...")
+    config = DeepXConfig()
+    pipeline = DeepXPipeline(config, embed_path="pretrained/gemma4_e2b_embed.pt")
+    
+    ckpt = torch.load(checkpoint_path, map_location="cpu")
+    sd = ckpt.get("model_state_dict", ckpt)
+    
+    # Dequantize if INT8
+    if ckpt.get("quantized", False):
+        logger.info("Dequantizing INT8 model...")
+        sd = dequantize_int8(sd)
+    
+    pipeline.backbone.load_state_dict(sd, strict=False)
+    pipeline = pipeline.to(device).half().eval()
+    
+    # Force pure GDN-2
+    for m in pipeline.modules():
+        if hasattr(m, 'path_mix_logit'):
+            m._alpha_override = 0.0
+    
+    logger.info("Model loaded successfully!")
+    logger.info(f"  VRAM: {torch.cuda.memory_allocated()/1024**2:.0f} MB" if device.type == "cuda" else "  CPU mode")
+
+
+@torch.no_grad()
+def encode_texts(texts: List[str], normalize: bool = True) -> torch.Tensor:
+    """Encode texts to embeddings."""
+    # Tokenize
+    encoded = tokenizer(
+        texts, padding=True, truncation=False, return_tensors="pt", max_length=MAX_LEN
+    )
+    
+    # Skip texts too long
+    input_ids = encoded["input_ids"]
+    attention_mask = encoded["attention_mask"]
+    
+    # Trim to max actual length (save compute)
+    max_len = attention_mask.sum(dim=1).max().item()
+    max_len = min(max_len, MAX_LEN)
+    input_ids = input_ids[:, :max_len].to(device)
+    attention_mask = attention_mask[:, :max_len].to(device)
+    
+    # Encode in batches
+    all_embs = []
+    for i in range(0, len(texts), MAX_BATCH):
+        batch_ids = input_ids[i:i+MAX_BATCH]
+        batch_mask = attention_mask[i:i+MAX_BATCH]
+        
+        with torch.amp.autocast(device_type="cuda", dtype=torch.float16):
+            emb = pipeline(batch_ids, attention_mask=batch_mask, normalize=normalize)
+        all_embs.append(emb.cpu())
+    
+    return torch.cat(all_embs)
+
+
+@app.route("/health", methods=["GET"])
+def health():
+    return jsonify({
+        "status": "ok",
+        "model": "deepx-v0.7-gdn2",
+        "device": str(device),
+        "max_len": MAX_LEN,
+        "dim": 1536,
+    })
+
+
+@app.route("/embed", methods=["POST"])
+def embed():
+    data = request.json
+    texts = data.get("texts", [])
+    normalize = data.get("normalize", True)
+    
+    if not texts:
+        return jsonify({"error": "No texts provided"}), 400
+    if len(texts) > 100:
+        return jsonify({"error": "Max 100 texts per request"}), 400
+    
+    t0 = time.time()
+    embeddings = encode_texts(texts, normalize=normalize)
+    elapsed = time.time() - t0
+    
+    return jsonify({
+        "embeddings": embeddings.tolist(),
+        "dim": embeddings.shape[1],
+        "count": len(texts),
+        "time_ms": round(elapsed * 1000, 1),
+    })
+
+
+@app.route("/similarity", methods=["POST"])
+def similarity():
+    data = request.json
+    query = data.get("query", "")
+    documents = data.get("documents", [])
+    
+    if not query or not documents:
+        return jsonify({"error": "Need query and documents"}), 400
+    
+    t0 = time.time()
+    all_texts = [query] + documents
+    embeddings = encode_texts(all_texts, normalize=True)
+    
+    query_emb = embeddings[0:1]
+    doc_embs = embeddings[1:]
+    scores = (query_emb @ doc_embs.T).squeeze(0).tolist()
+    elapsed = time.time() - t0
+    
+    return jsonify({
+        "scores": scores,
+        "time_ms": round(elapsed * 1000, 1),
+    })
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--checkpoint", default="checkpoints/deploy/deepx_int8.pt")
+    parser.add_argument("--tokenizer", default="/mnt/d/DX/hf_cache/hub/models--google--gemma-4-E2B-it/snapshots/6b7e72c67d3c4556f42b56d5a68b4b8e864c63b4")
+    parser.add_argument("--port", type=int, default=8080)
+    parser.add_argument("--host", default="0.0.0.0")
+    parser.add_argument("--device", default="cuda")
+    args = parser.parse_args()
+    
+    load_model(args.checkpoint, args.tokenizer, args.device)
+    
+    logger.info(f"Starting server on {args.host}:{args.port}")
+    app.run(host=args.host, port=args.port, debug=False)
+
+
+if __name__ == "__main__":
+    main()