specmem-hardwicksoftware 3.5.99

package/scripts/optimize-embedding-model.py

@@ -0,0 +1,324 @@
+#!/usr/bin/env python3
+"""
+SPECMEM EMBEDDING MODEL OPTIMIZER
+=================================
+
+Optimizes the embedding model for CPU inference WITHOUT GPU!
+- ONNX export
+- INT8 quantization (2-4x speedup)
+- Benchmark before/after
+
+Run on server with 32GB RAM - NO GPU NEEDED!
+
+@author hardwicksoftwareservices
+@website https://justcalljon.pro
+"""
+
+import os
+import sys
+import time
+import json
+import hashlib
+from pathlib import Path
+
+print("""
+╔═══════════════════════════════════════════════════════════════╗
+║   SPECMEM EMBEDDING MODEL OPTIMIZER                           ║
+║   CPU-Only Optimization (No GPU Required!)                    ║
+║   https://justcalljon.pro                                     ║
+╚═══════════════════════════════════════════════════════════════╝
+""")
+
+# Output directory
+OUTPUT_DIR = Path("/specmem/models/optimized")
+OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+
+# Check dependencies
+print("═══ Checking Dependencies ═══\n")
+
+try:
+    import torch
+    print(f"✓ PyTorch {torch.__version__}")
+except ImportError:
+    print("✗ PyTorch not found - installing...")
+    os.system("pip install torch --index-url https://download.pytorch.org/whl/cpu")
+    import torch
+
+try:
+    from sentence_transformers import SentenceTransformer
+    print(f"✓ sentence-transformers installed")
+except ImportError:
+    print("✗ sentence-transformers not found - installing...")
+    os.system("pip install sentence-transformers[onnx]")
+    from sentence_transformers import SentenceTransformer
+
+try:
+    import onnx
+    print(f"✓ ONNX {onnx.__version__}")
+except ImportError:
+    print("✗ ONNX not found - installing...")
+    os.system("pip install onnx onnxruntime")
+    import onnx
+
+try:
+    from optimum.onnxruntime import ORTModelForFeatureExtraction
+    from optimum.onnxruntime.configuration import AutoQuantizationConfig
+    from optimum.onnxruntime import ORTQuantizer
+    print(f"✓ Hugging Face Optimum installed")
+except ImportError:
+    print("✗ Optimum not found - installing...")
+    os.system("pip install optimum[onnxruntime]")
+    from optimum.onnxruntime import ORTModelForFeatureExtraction
+    from optimum.onnxruntime.configuration import AutoQuantizationConfig
+    from optimum.onnxruntime import ORTQuantizer
+
+import numpy as np
+
+# Model to optimize
+MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
+
+print(f"\n═══ Loading Base Model ═══\n")
+print(f"Model: {MODEL_NAME}")
+
+# Load the base PyTorch model for benchmarking
+print("Loading PyTorch model...")
+start = time.time()
+pytorch_model = SentenceTransformer(MODEL_NAME)
+pytorch_load_time = time.time() - start
+print(f"✓ PyTorch model loaded in {pytorch_load_time:.2f}s")
+
+# Test sentences for benchmarking
+TEST_SENTENCES = [
+    "The quick brown fox jumps over the lazy dog.",
+    "Machine learning models can be optimized for CPU inference.",
+    "SpecMem provides semantic memory for  Code.",
+    "Quantization reduces model size and increases inference speed.",
+    "PostgreSQL with pgvector enables efficient vector similarity search.",
+] * 20  # 100 sentences total
+
+print(f"\n═══ Benchmarking PyTorch Model ═══\n")
+print(f"Running {len(TEST_SENTENCES)} embeddings...")
+
+# Warm up
+_ = pytorch_model.encode(TEST_SENTENCES[:5])
+
+# Benchmark PyTorch
+start = time.time()
+pytorch_embeddings = pytorch_model.encode(TEST_SENTENCES)
+pytorch_time = time.time() - start
+pytorch_per_sentence = (pytorch_time / len(TEST_SENTENCES)) * 1000
+
+print(f"✓ PyTorch: {pytorch_time:.3f}s total, {pytorch_per_sentence:.2f}ms per sentence")
+
+# Get model size
+pytorch_size = sum(p.numel() * p.element_size() for p in pytorch_model[0].auto_model.parameters())
+print(f"  Model size: {pytorch_size / 1024 / 1024:.2f} MB")
+
+print(f"\n═══ Exporting to ONNX ═══\n")
+
+onnx_path = OUTPUT_DIR / "model.onnx"
+onnx_quantized_path = OUTPUT_DIR / "model_quantized.onnx"
+
+# Export to ONNX using sentence-transformers
+print("Exporting to ONNX format...")
+start = time.time()
+
+try:
+    # Try using sentence-transformers ONNX backend
+    onnx_model = SentenceTransformer(MODEL_NAME, backend="onnx")
+
+    # Save the ONNX model
+    onnx_model.save(str(OUTPUT_DIR / "onnx_model"))
+    export_time = time.time() - start
+    print(f"✓ ONNX export completed in {export_time:.2f}s")
+except Exception as e:
+    print(f"  Using alternative export method: {e}")
+    # Alternative: use optimum
+    from optimum.onnxruntime import ORTModelForFeatureExtraction
+    from transformers import AutoTokenizer
+
+    ort_model = ORTModelForFeatureExtraction.from_pretrained(MODEL_NAME, export=True)
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+
+    ort_model.save_pretrained(str(OUTPUT_DIR / "onnx_model"))
+    tokenizer.save_pretrained(str(OUTPUT_DIR / "onnx_model"))
+    export_time = time.time() - start
+    print(f"✓ ONNX export (via Optimum) completed in {export_time:.2f}s")
+
+print(f"\n═══ Quantizing to INT8 ═══\n")
+print("This is where the MAGIC happens - 2-4x speedup!")
+
+try:
+    # Try sentence-transformers quantization
+    from sentence_transformers import export_dynamic_quantized_onnx_model
+
+    print("Using sentence-transformers quantization...")
+    start = time.time()
+
+    # Reload as ONNX model
+    onnx_model = SentenceTransformer(MODEL_NAME, backend="onnx")
+
+    # Export quantized version
+    export_dynamic_quantized_onnx_model(
+        model=onnx_model,
+        quantization_config="avx512",  # Use AVX512 for modern CPUs
+        model_name_or_path=str(OUTPUT_DIR / "quantized_model")
+    )
+
+    quant_time = time.time() - start
+    print(f"✓ INT8 quantization completed in {quant_time:.2f}s")
+
+except Exception as e:
+    print(f"  Using Optimum quantization: {e}")
+
+    from optimum.onnxruntime import ORTQuantizer
+    from optimum.onnxruntime.configuration import AutoQuantizationConfig
+
+    # Load the ONNX model
+    quantizer = ORTQuantizer.from_pretrained(str(OUTPUT_DIR / "onnx_model"))
+
+    # Configure quantization
+    qconfig = AutoQuantizationConfig.avx512_vnni(is_static=False, per_channel=False)
+
+    # Quantize
+    start = time.time()
+    quantizer.quantize(
+        save_dir=str(OUTPUT_DIR / "quantized_model"),
+        quantization_config=qconfig,
+    )
+    quant_time = time.time() - start
+    print(f"✓ INT8 quantization (via Optimum) completed in {quant_time:.2f}s")
+
+print(f"\n═══ Benchmarking Optimized Models ═══\n")
+
+# Benchmark ONNX model
+print("Loading ONNX model...")
+try:
+    onnx_model = SentenceTransformer(str(OUTPUT_DIR / "onnx_model"), backend="onnx")
+
+    # Warm up
+    _ = onnx_model.encode(TEST_SENTENCES[:5])
+
+    # Benchmark
+    start = time.time()
+    onnx_embeddings = onnx_model.encode(TEST_SENTENCES)
+    onnx_time = time.time() - start
+    onnx_per_sentence = (onnx_time / len(TEST_SENTENCES)) * 1000
+
+    print(f"✓ ONNX: {onnx_time:.3f}s total, {onnx_per_sentence:.2f}ms per sentence")
+    print(f"  Speedup vs PyTorch: {pytorch_time/onnx_time:.2f}x")
+except Exception as e:
+    print(f"  ONNX benchmark failed: {e}")
+    onnx_time = pytorch_time
+    onnx_per_sentence = pytorch_per_sentence
+
+# Benchmark quantized model
+print("\nLoading INT8 quantized model...")
+try:
+    # Find the quantized model file
+    quant_model_path = OUTPUT_DIR / "quantized_model"
+
+    quantized_model = SentenceTransformer(
+        str(quant_model_path),
+        backend="onnx",
+        model_kwargs={"file_name": "model_quantized.onnx"}
+    )
+
+    # Warm up
+    _ = quantized_model.encode(TEST_SENTENCES[:5])
+
+    # Benchmark
+    start = time.time()
+    quantized_embeddings = quantized_model.encode(TEST_SENTENCES)
+    quantized_time = time.time() - start
+    quantized_per_sentence = (quantized_time / len(TEST_SENTENCES)) * 1000
+
+    print(f"✓ INT8 Quantized: {quantized_time:.3f}s total, {quantized_per_sentence:.2f}ms per sentence")
+    print(f"  Speedup vs PyTorch: {pytorch_time/quantized_time:.2f}x")
+    print(f"  Speedup vs ONNX: {onnx_time/quantized_time:.2f}x")
+except Exception as e:
+    print(f"  Quantized benchmark skipped: {e}")
+    quantized_time = onnx_time
+    quantized_per_sentence = onnx_per_sentence
+
+# Calculate file sizes
+print(f"\n═══ Model Sizes ═══\n")
+
+def get_dir_size(path):
+    total = 0
+    for f in Path(path).rglob("*"):
+        if f.is_file():
+            total += f.stat().st_size
+    return total
+
+onnx_size = get_dir_size(OUTPUT_DIR / "onnx_model") if (OUTPUT_DIR / "onnx_model").exists() else 0
+quant_size = get_dir_size(OUTPUT_DIR / "quantized_model") if (OUTPUT_DIR / "quantized_model").exists() else 0
+
+print(f"PyTorch model:    {pytorch_size / 1024 / 1024:.2f} MB (in memory)")
+print(f"ONNX model:       {onnx_size / 1024 / 1024:.2f} MB (on disk)")
+print(f"Quantized model:  {quant_size / 1024 / 1024:.2f} MB (on disk)")
+
+if quant_size > 0 and onnx_size > 0:
+    print(f"Size reduction:   {(1 - quant_size/onnx_size) * 100:.1f}%")
+
+# Verify embedding quality
+print(f"\n═══ Quality Check ═══\n")
+
+# Compare embeddings similarity
+from numpy.linalg import norm
+
+def cosine_similarity(a, b):
+    return np.dot(a, b) / (norm(a) * norm(b))
+
+if 'onnx_embeddings' in dir() and 'pytorch_embeddings' in dir():
+    sim = cosine_similarity(pytorch_embeddings[0], onnx_embeddings[0])
+    print(f"PyTorch vs ONNX similarity: {sim:.6f}")
+
+if 'quantized_embeddings' in dir() and 'pytorch_embeddings' in dir():
+    sim = cosine_similarity(pytorch_embeddings[0], quantized_embeddings[0])
+    print(f"PyTorch vs INT8 similarity: {sim:.6f}")
+
+print("(Values > 0.99 = excellent quality preservation)")
+
+# Summary
+print(f"""
+╔═══════════════════════════════════════════════════════════════╗
+║                      OPTIMIZATION SUMMARY                      ║
+╠═══════════════════════════════════════════════════════════════╣
+║  Model: {MODEL_NAME:<43} ║
+╠═══════════════════════════════════════════════════════════════╣
+║  BENCHMARK RESULTS ({len(TEST_SENTENCES)} sentences)                           ║
+║  ─────────────────────────────────────────────────────────────║
+║  PyTorch:     {pytorch_per_sentence:>6.2f} ms/sentence                          ║
+║  ONNX:        {onnx_per_sentence:>6.2f} ms/sentence  ({pytorch_time/onnx_time:.1f}x faster)              ║
+║  INT8:        {quantized_per_sentence:>6.2f} ms/sentence  ({pytorch_time/quantized_time:.1f}x faster)              ║
+╠═══════════════════════════════════════════════════════════════╣
+║  Output: {str(OUTPUT_DIR):<50} ║
+╚═══════════════════════════════════════════════════════════════╝
+""")
+
+# Create manifest for hosting
+manifest = {
+    "version": "1.0.0",
+    "model_name": "specmem-embedding-v1",
+    "base_model": MODEL_NAME,
+    "optimizations": ["onnx", "int8_quantization"],
+    "embedding_dim": 384,
+    "max_seq_length": 256,
+    "benchmark": {
+        "pytorch_ms": round(pytorch_per_sentence, 2),
+        "onnx_ms": round(onnx_per_sentence, 2),
+        "quantized_ms": round(quantized_per_sentence, 2),
+        "speedup": round(pytorch_time/quantized_time, 2)
+    },
+    "files": {
+        "onnx_model": "onnx_model/",
+        "quantized_model": "quantized_model/"
+    }
+}
+
+with open(OUTPUT_DIR / "manifest.json", "w") as f:
+    json.dump(manifest, f, indent=2)
+
+print(f"✓ Manifest saved to {OUTPUT_DIR / 'manifest.json'}")
+print(f"\n🚀 Models ready for hosting on justcalljon.pro!")