trying hibrid approach

app.py

@@ -1,5 +1,6 @@
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
+from contextlib import asynccontextmanager
 from typing import List
 import torch
 import uvicorn
@@ -7,10 +8,53 @@ import uvicorn
 from models.schemas import EmbeddingRequest, EmbeddingResponse, ModelInfo
 from utils.helpers import load_models, get_embeddings, cleanup_memory
 
+# Global model cache
+models_cache = {}
+
+# Models to load at startup (most frequently used)
+STARTUP_MODELS = ["jina-v3", "roberta-ca"]
+# Models to load on demand
+ON_DEMAND_MODELS = ["jina", "robertalex", "legal-bert"]
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Application lifespan handler for startup and shutdown"""
+    # Startup - load priority models
+    try:
+        global models_cache
+        print(f"Loading startup models: {STARTUP_MODELS}...")
+        models_cache = load_models(STARTUP_MODELS)
+        print(f"Startup models loaded successfully: {list(models_cache.keys())}")
+        yield
+    except Exception as e:
+        print(f"Failed to load startup models: {str(e)}")
+        # Continue anyway - models can be loaded on demand
+        yield
+    finally:
+        # Shutdown - cleanup resources
+        cleanup_memory()
+
+def ensure_model_loaded(model_name: str):
+    """Load a specific model on demand if not already loaded"""
+    global models_cache
+    if model_name not in models_cache:
+        if model_name in ON_DEMAND_MODELS:
+            try:
+                print(f"Loading model on demand: {model_name}...")
+                new_models = load_models([model_name])
+                models_cache.update(new_models)
+                print(f"Model {model_name} loaded successfully!")
+            except Exception as e:
+                print(f"Failed to load model {model_name}: {str(e)}")
+                raise HTTPException(status_code=500, detail=f"Model {model_name} loading failed: {str(e)}")
+        else:
+            raise HTTPException(status_code=400, detail=f"Unknown model: {model_name}")
+
 app = FastAPI(
     title="Multilingual & Legal Embedding API",
     description="Multi-model embedding API for Spanish, Catalan, English and Legal texts",
-    version="3.0.0"
+    version="3.0.0",
+    lifespan=lifespan
 )
 
 # Add CORS middleware to allow cross-origin requests
@@ -22,21 +66,6 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Global model cache - loaded on demand
-models_cache = {}
-
-def ensure_models_loaded():
-    """Load models on first request if not already loaded"""
-    global models_cache
-    if not models_cache:
-        try:
-            print("Loading models on demand...")
-            models_cache = load_models()
-            print("All models loaded successfully!")
-        except Exception as e:
-            print(f"Failed to load models: {str(e)}")
-            raise HTTPException(status_code=500, detail=f"Model loading failed: {str(e)}")
-
 @app.get("/")
 async def root():
     return {
@@ -51,8 +80,8 @@ async def root():
 async def create_embeddings(request: EmbeddingRequest):
     """Generate embeddings for input texts"""
     try:
-        # Load models on first request
-        ensure_models_loaded()
+        # Load specific model on demand if needed
+        ensure_model_loaded(request.model)
         
         if not request.texts:
             raise HTTPException(status_code=400, detail="No texts provided")
@@ -138,14 +167,18 @@ async def list_models():
 @app.get("/health")
 async def health_check():
     """Health check endpoint"""
-    models_loaded = len(models_cache) == 5
+    startup_models_loaded = all(model in models_cache for model in STARTUP_MODELS)
+    all_models_loaded = len(models_cache) == 5
+    
     return {
-        "status": "healthy" if models_loaded else "ready",
-        "models_loaded": models_loaded,
+        "status": "healthy" if startup_models_loaded else "partial",
+        "startup_models_loaded": startup_models_loaded,
+        "all_models_loaded": all_models_loaded,
         "available_models": list(models_cache.keys()),
-        "expected_models": ["jina", "robertalex", "jina-v3", "legal-bert", "roberta-ca"],
+        "startup_models": STARTUP_MODELS,
+        "on_demand_models": ON_DEMAND_MODELS,
         "models_count": len(models_cache),
-        "note": "Models load on first embedding request" if not models_loaded else "All models ready"
+        "note": f"Startup models: {STARTUP_MODELS} | On-demand: {ON_DEMAND_MODELS}"
     }
 
 if __name__ == "__main__":

test_hybrid.py

@@ -0,0 +1,98 @@
+#!/usr/bin/env python3
+"""
+Test script for hybrid model loading
+"""
+
+import requests
+import json
+import time
+
+def test_hybrid_api(base_url="https://aurasystems-spanish-embeddings-api.hf.space"):
+    """Test the hybrid API"""
+    
+    print(f"Testing hybrid API at {base_url}")
+    
+    # Test health endpoint first
+    try:
+        response = requests.get(f"{base_url}/health")
+        print(f"✓ Health endpoint: {response.status_code}")
+        if response.status_code == 200:
+            health_data = response.json()
+            print(f"  Startup models loaded: {health_data.get('startup_models_loaded', False)}")
+            print(f"  Available models: {health_data.get('available_models', [])}")
+            print(f"  Note: {health_data.get('note', 'N/A')}")
+        else:
+            print(f"  Error: {response.text}")
+    except Exception as e:
+        print(f"✗ Health endpoint failed: {e}")
+        return False
+    
+    # Test startup model (jina-v3)
+    try:
+        payload = {
+            "texts": ["Hola mundo", "Bonjour le monde"],
+            "model": "jina-v3",
+            "normalize": True
+        }
+        response = requests.post(f"{base_url}/embed", json=payload)
+        print(f"✓ Startup model (jina-v3): {response.status_code}")
+        if response.status_code == 200:
+            data = response.json()
+            print(f"  Generated {data.get('num_texts', 0)} embeddings")
+            print(f"  Dimensions: {data.get('dimensions', 0)}")
+        else:
+            print(f"  Error: {response.text}")
+    except Exception as e:
+        print(f"✗ Startup model test failed: {e}")
+    
+    # Test startup model (roberta-ca)
+    try:
+        payload = {
+            "texts": ["Bon dia", "Com estàs?"],
+            "model": "roberta-ca",
+            "normalize": True
+        }
+        response = requests.post(f"{base_url}/embed", json=payload)
+        print(f"✓ Startup model (roberta-ca): {response.status_code}")
+        if response.status_code == 200:
+            data = response.json()
+            print(f"  Generated {data.get('num_texts', 0)} embeddings")
+            print(f"  Dimensions: {data.get('dimensions', 0)}")
+        else:
+            print(f"  Error: {response.text}")
+    except Exception as e:
+        print(f"✗ Startup model test failed: {e}")
+    
+    # Test on-demand model (jina)
+    try:
+        payload = {
+            "texts": ["Texto en español"],
+            "model": "jina",
+            "normalize": True
+        }
+        response = requests.post(f"{base_url}/embed", json=payload)
+        print(f"✓ On-demand model (jina): {response.status_code}")
+        if response.status_code == 200:
+            data = response.json()
+            print(f"  Generated {data.get('num_texts', 0)} embeddings")
+            print(f"  Dimensions: {data.get('dimensions', 0)}")
+        else:
+            print(f"  Error: {response.text}")
+    except Exception as e:
+        print(f"✗ On-demand model test failed: {e}")
+    
+    # Check health again to see all models
+    try:
+        response = requests.get(f"{base_url}/health")
+        if response.status_code == 200:
+            health_data = response.json()
+            print(f"✓ Final health check:")
+            print(f"  All models loaded: {health_data.get('all_models_loaded', False)}")
+            print(f"  Available models: {health_data.get('available_models', [])}")
+    except Exception as e:
+        print(f"✗ Final health check failed: {e}")
+    
+    return True
+
+if __name__ == "__main__":
+    test_hybrid_api()

utils/helpers.py

@@ -12,104 +12,118 @@ from typing import List, Dict, Optional
 import gc
 import os
 
-def load_models() -> Dict:
+def load_models(model_names: List[str] = None) -> Dict:
     """
-    Load all embedding models with memory optimization
+    Load specific embedding models with memory optimization
+    
+    Args:
+        model_names: List of model names to load. If None, loads all models.
     
     Returns:
         Dict containing loaded models and tokenizers
     """
     models_cache = {}
     
+    # Default to all models if none specified
+    if model_names is None:
+        model_names = ["jina", "robertalex", "jina-v3", "legal-bert", "roberta-ca"]
+    
     # Set device
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     
     try:
         # Load Jina v2 Spanish model
-        print("Loading Jina embeddings v2 Spanish model...")
-        jina_tokenizer = AutoTokenizer.from_pretrained(
-            'jinaai/jina-embeddings-v2-base-es',
-            trust_remote_code=True
-        )
-        jina_model = AutoModel.from_pretrained(
-            'jinaai/jina-embeddings-v2-base-es',
-            trust_remote_code=True,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-        ).to(device)
-        jina_model.eval()
-        
-        # Load RoBERTalex model
-        print("Loading RoBERTalex model...")
-        robertalex_tokenizer = RobertaTokenizer.from_pretrained('PlanTL-GOB-ES/RoBERTalex')
-        robertalex_model = RobertaModel.from_pretrained(
-            'PlanTL-GOB-ES/RoBERTalex',
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-        ).to(device)
-        robertalex_model.eval()
-        
-        # Load Jina v3 model
-        print("Loading Jina embeddings v3 model...")
-        jina_v3_tokenizer = AutoTokenizer.from_pretrained(
-            'jinaai/jina-embeddings-v3',
-            trust_remote_code=True
-        )
-        jina_v3_model = AutoModel.from_pretrained(
-            'jinaai/jina-embeddings-v3',
-            trust_remote_code=True,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-        ).to(device)
-        jina_v3_model.eval()
-        
-        # Load Legal BERT model
-        print("Loading Legal BERT model...")
-        legal_bert_tokenizer = BertTokenizer.from_pretrained('nlpaueb/legal-bert-base-uncased')
-        legal_bert_model = BertModel.from_pretrained(
-            'nlpaueb/legal-bert-base-uncased',
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-        ).to(device)
-        legal_bert_model.eval()
-        
-        # Load Catalan RoBERTa model
-        print("Loading Catalan RoBERTa-large model...")
-        roberta_ca_tokenizer = AutoTokenizer.from_pretrained('projecte-aina/roberta-large-ca-v2')
-        roberta_ca_model = AutoModel.from_pretrained(
-            'projecte-aina/roberta-large-ca-v2',
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-        ).to(device)
-        roberta_ca_model.eval()
-        
-        models_cache = {
-            'jina': {
+        if "jina" in model_names:
+            print("Loading Jina embeddings v2 Spanish model...")
+            jina_tokenizer = AutoTokenizer.from_pretrained(
+                'jinaai/jina-embeddings-v2-base-es',
+                trust_remote_code=True
+            )
+            jina_model = AutoModel.from_pretrained(
+                'jinaai/jina-embeddings-v2-base-es',
+                trust_remote_code=True,
+                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+            ).to(device)
+            jina_model.eval()
+            
+            models_cache['jina'] = {
                 'tokenizer': jina_tokenizer,
                 'model': jina_model,
                 'device': device,
                 'pooling': 'mean'
-            },
-            'robertalex': {
+            }
+        
+        # Load RoBERTalex model
+        if "robertalex" in model_names:
+            print("Loading RoBERTalex model...")
+            robertalex_tokenizer = RobertaTokenizer.from_pretrained('PlanTL-GOB-ES/RoBERTalex')
+            robertalex_model = RobertaModel.from_pretrained(
+                'PlanTL-GOB-ES/RoBERTalex',
+                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+            ).to(device)
+            robertalex_model.eval()
+            
+            models_cache['robertalex'] = {
                 'tokenizer': robertalex_tokenizer,
                 'model': robertalex_model,
                 'device': device,
                 'pooling': 'cls'
-            },
-            'jina-v3': {
+            }
+        
+        # Load Jina v3 model
+        if "jina-v3" in model_names:
+            print("Loading Jina embeddings v3 model...")
+            jina_v3_tokenizer = AutoTokenizer.from_pretrained(
+                'jinaai/jina-embeddings-v3',
+                trust_remote_code=True
+            )
+            jina_v3_model = AutoModel.from_pretrained(
+                'jinaai/jina-embeddings-v3',
+                trust_remote_code=True,
+                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+            ).to(device)
+            jina_v3_model.eval()
+            
+            models_cache['jina-v3'] = {
                 'tokenizer': jina_v3_tokenizer,
                 'model': jina_v3_model,
                 'device': device,
                 'pooling': 'mean'
-            },
-            'legal-bert': {
+            }
+        
+        # Load Legal BERT model
+        if "legal-bert" in model_names:
+            print("Loading Legal BERT model...")
+            legal_bert_tokenizer = BertTokenizer.from_pretrained('nlpaueb/legal-bert-base-uncased')
+            legal_bert_model = BertModel.from_pretrained(
+                'nlpaueb/legal-bert-base-uncased',
+                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+            ).to(device)
+            legal_bert_model.eval()
+            
+            models_cache['legal-bert'] = {
                 'tokenizer': legal_bert_tokenizer,
                 'model': legal_bert_model,
                 'device': device,
                 'pooling': 'cls'
-            },
-            'roberta-ca': {
+            }
+        
+        # Load Catalan RoBERTa model
+        if "roberta-ca" in model_names:
+            print("Loading Catalan RoBERTa-large model...")
+            roberta_ca_tokenizer = AutoTokenizer.from_pretrained('projecte-aina/roberta-large-ca-v2')
+            roberta_ca_model = AutoModel.from_pretrained(
+                'projecte-aina/roberta-large-ca-v2',
+                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+            ).to(device)
+            roberta_ca_model.eval()
+            
+            models_cache['roberta-ca'] = {
                 'tokenizer': roberta_ca_tokenizer,
                 'model': roberta_ca_model,
                 'device': device,
                 'pooling': 'cls'
             }
-        }
         
         # Force garbage collection after loading
         gc.collect()