added caches

requirements.txt

@@ -7,4 +7,6 @@ sacremoses
 torch
 pillow
 protobuf
-# Optional dependencies for specific features
+
+# Optional dependencies for specific features
+einops

src/embeddings.py

@@ -62,7 +62,7 @@ class BaseEmbeddingTaskService:
         """Load and cache processor for the model using AutoProcessor"""
         if model_name not in self._processor_cache:
             try:
-                self._processor_cache[model_name] = AutoProcessor.from_pretrained(model_name)
+                self._processor_cache[model_name] = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
                 self._logger.info(f"Loaded processor for model: {model_name}")
             except Exception as e:
                 self._logger.error(f"Failed to load processor for model '{model_name}': {str(e)}")
@@ -70,6 +70,8 @@ class BaseEmbeddingTaskService:
                     status_code=404,
                     detail=f"Processor for model '{model_name}' could not be loaded: {str(e)}"
                 )
+        else:
+            self._logger.info(f"Using cached processor for model: {model_name}")
         return self._processor_cache[model_name]
 
     def _load_model(self, model_name: str, cache_suffix: str = ""):
@@ -78,7 +80,7 @@ class BaseEmbeddingTaskService:
         if cache_key not in self._model_cache:
             try:
                 device = self._get_device()
-                model = AutoModel.from_pretrained(model_name)
+                model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
                 model.to(device)
                 self._model_cache[cache_key] = model
                 self._logger.info(f"Loaded model: {model_name} on {device}")
@@ -88,6 +90,8 @@ class BaseEmbeddingTaskService:
                     status_code=404,
                     detail=f"Model '{model_name}' could not be loaded: {str(e)}"
                 )
+        else:
+            self._logger.info(f"Using cached model: {model_name} (cache key: {cache_key})")
         return self._model_cache[cache_key]
 
     async def get_embedding_vector_size(self, model_name: str) -> dict:
@@ -335,7 +339,7 @@ class TextEmbeddingTaskService(BaseEmbeddingTaskService):
         """Main method to generate text embeddings"""
         embedding_request: EmbeddingRequest = await self.get_embedding_request(request)
         
-        self._logger.info(f"Generating text embedding for: {embedding_request.inputs[:50]}...")
+        self._logger.info(f"Generating text embedding for: {embedding_request.inputs[:500]}...")
         
         # Load processor and model using auto-detection
         processor = self._load_processor(model_name)

src/main.py

@@ -29,6 +29,10 @@ logging.basicConfig(format='%(asctime)s %(levelname)-8s %(message)s')
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
+# Create singleton instances of embedding services to enable model caching across requests
+image_embedding_service = ImageEmbeddingTaskService(logger)
+text_embedding_service = TextEmbeddingTaskService(logger)
+
 
 class StreamToLogger(object):
     def __init__(self, logger, log_level):
@@ -333,8 +337,7 @@ async def image_embedding(
     """
 
     model_name = model_name.rstrip("/")
-    imageEmbeddingTask = ImageEmbeddingTaskService(logger)
-    return await imageEmbeddingTask.generate_embedding(request, model_name)
+    return await image_embedding_service.generate_embedding(request, model_name)
 
 
 # =========================
@@ -399,8 +402,7 @@ async def image_embedding_upload(
     """
 
     model_name = model_name.rstrip("/")
-    imageEmbeddingTask = ImageEmbeddingTaskService(logger)
-    return await imageEmbeddingTask.generate_embedding_from_upload(image, model_name)
+    return await image_embedding_service.generate_embedding_from_upload(image, model_name)
 
 
 # =========================
@@ -439,8 +441,7 @@ async def text_embedding(
     """
 
     model_name = model_name.rstrip("/")
-    textEmbeddingTask = TextEmbeddingTaskService(logger)
-    return await textEmbeddingTask.generate_embedding(request, model_name)
+    return await text_embedding_service.generate_embedding(request, model_name)
 
 
 # =========================
@@ -483,6 +484,5 @@ async def embedding_vector_size(
     """
 
     model_name = model_name.rstrip("/")
-    # We can use either ImageEmbeddingTaskService or TextEmbeddingTaskService as they inherit from the same base class
-    embeddingTask = ImageEmbeddingTaskService(logger)
-    return await embeddingTask.get_embedding_vector_size(model_name)
+    # We can use either embedding service as they inherit from the same base class
+    return await image_embedding_service.get_embedding_vector_size(model_name)