Added double init, for embedding and chat models at the same time.

main/api.py

@@ -17,9 +17,13 @@ class LLMApi:
         self.models_path = self.base_path / config["folders"]["models"]
         self.cache_path = self.base_path / config["folders"]["cache"]
 
-        self.model = None
-        self.model_name = None
+        # Initialize model variables for both generation and embedding
+        self.generation_model = None
+        self.generation_model_name = None
+        self.embedding_model = None
+        self.embedding_model_name = None
         self.tokenizer = None
+        self.embedding_tokenizer = None
 
         # Generation parameters from config
         gen_config = config["model"]["generation"]
@@ -64,7 +68,7 @@ class LLMApi:
 
             # Download and save tokenizer
             tokenizer = AutoTokenizer.from_pretrained(model_name)
-            self.logger.info(f"Disnabling stdout logging")
+            self.logger.info(f"Disabling stdout logging")
             self.logger.disable_stream_to_logger()
 
             self.logger.info(f"Saving model to {model_path}")
@@ -78,14 +82,14 @@ class LLMApi:
 
     def initialize_model(self, model_name: str) -> None:
         """
-        Initialize a model and tokenizer, either from local storage or by downloading.
+        Initialize a model and tokenizer for text generation.
 
         Args:
             model_name: The name of the model to initialize
         """
-        self.logger.info(f"Initializing model: {model_name}")
+        self.logger.info(f"Initializing generation model: {model_name}")
         try:
-            self.model_name = model_name
+            self.generation_model_name = model_name
             local_model_path = self.models_path / model_name.split('/')[-1]
 
             # Check if model exists locally
@@ -96,7 +100,7 @@ class LLMApi:
                 self.logger.info(f"Loading model from source: {model_name}")
                 model_path = model_name
 
-            self.model = AutoModelForCausalLM.from_pretrained(
+            self.generation_model = AutoModelForCausalLM.from_pretrained(
                 model_path,
                 device_map="auto",
                 load_in_8bit=True,
@@ -108,9 +112,42 @@ class LLMApi:
             self.generation_config["eos_token_id"] = self.tokenizer.eos_token_id
             self.generation_config["pad_token_id"] = self.tokenizer.eos_token_id
 
-            self.logger.info(f"Successfully initialized model: {model_name}")
+            self.logger.info(f"Successfully initialized generation model: {model_name}")
         except Exception as e:
-            self.logger.error(f"Failed to initialize model {model_name}: {str(e)}")
+            self.logger.error(f"Failed to initialize generation model {model_name}: {str(e)}")
+            raise
+
+    def initialize_embedding_model(self, model_name: str) -> None:
+        """
+        Initialize a model and tokenizer specifically for embeddings.
+
+        Args:
+            model_name: The name of the model to initialize for embeddings
+        """
+        self.logger.info(f"Initializing embedding model: {model_name}")
+        try:
+            self.embedding_model_name = model_name
+            local_model_path = self.models_path / model_name.split('/')[-1]
+
+            # Check if model exists locally
+            if local_model_path.exists():
+                self.logger.info(f"Loading embedding model from local path: {local_model_path}")
+                model_path = local_model_path
+            else:
+                self.logger.info(f"Loading embedding model from source: {model_name}")
+                model_path = model_name
+
+            self.embedding_model = AutoModelForCausalLM.from_pretrained(
+                model_path,
+                device_map="auto",
+                load_in_8bit=True,
+                torch_dtype=torch.float16
+            )
+            self.embedding_tokenizer = AutoTokenizer.from_pretrained(model_path)
+
+            self.logger.info(f"Successfully initialized embedding model: {model_name}")
+        except Exception as e:
+            self.logger.error(f"Failed to initialize embedding model {model_name}: {str(e)}")
             raise
 
     def has_chat_template(self) -> bool:
@@ -158,22 +195,22 @@ class LLMApi:
         """
         self.logger.debug(f"Generating response for prompt: {prompt[:50]}...")
 
-        if self.model is None:
-            raise RuntimeError("Model not initialized. Call initialize_model first.")
+        if self.generation_model is None:
+            raise RuntimeError("Generation model not initialized. Call initialize_model first.")
 
         try:
             text = self._prepare_prompt(prompt, system_message)
             inputs = self.tokenizer([text], return_tensors="pt")
 
             # Remove token_type_ids if present
-            model_inputs = {k: v.to(self.model.device) for k, v in inputs.items()
+            model_inputs = {k: v.to(self.generation_model.device) for k, v in inputs.items()
                             if k != 'token_type_ids'}
 
             generation_config = self.generation_config.copy()
             if max_new_tokens:
                 generation_config["max_new_tokens"] = max_new_tokens
 
-            generated_ids = self.model.generate(
+            generated_ids = self.generation_model.generate(
                 **model_inputs,
                 **generation_config
             )
@@ -202,15 +239,15 @@ class LLMApi:
         """
         self.logger.debug(f"Starting streaming generation for prompt: {prompt[:50]}...")
 
-        if self.model is None:
-            raise RuntimeError("Model not initialized. Call initialize_model first.")
+        if self.generation_model is None:
+            raise RuntimeError("Generation model not initialized. Call initialize_model first.")
 
         try:
             text = self._prepare_prompt(prompt, system_message)
             inputs = self.tokenizer([text], return_tensors="pt")
 
             # Remove token_type_ids if present
-            model_inputs = {k: v.to(self.model.device) for k, v in inputs.items()
+            model_inputs = {k: v.to(self.generation_model.device) for k, v in inputs.items()
                             if k != 'token_type_ids'}
 
             # Configure generation
@@ -227,7 +264,7 @@ class LLMApi:
             )
 
             # Create a thread to run the generation
-            thread = Thread(target=self.model.generate, kwargs=generation_kwargs)
+            thread = Thread(target=self.generation_model.generate, kwargs=generation_kwargs)
             thread.start()
 
             # Yield the generated text in chunks
@@ -241,21 +278,21 @@ class LLMApi:
 
     def generate_embedding(self, text: str) -> List[float]:
         """
-        Generate a single embedding vector for a chunk of text.
+        Generate a single embedding vector for a chunk of text using the dedicated embedding model.
         Returns a list of floats representing the text embedding.
         """
         self.logger.debug(f"Generating embedding for text: {text[:50]}...")
 
-        if self.model is None or self.tokenizer is None:
-            raise RuntimeError("Model not initialized. Call initialize_model first.")
+        if self.embedding_model is None or self.embedding_tokenizer is None:
+            raise RuntimeError("Embedding model not initialized. Call initialize_embedding_model first.")
 
         try:
             # Tokenize the input text and ensure input_ids are Long type
-            inputs = self.tokenizer(text, return_tensors='pt')
-            input_ids = inputs.input_ids.to(dtype=torch.long, device=self.model.device)
+            inputs = self.embedding_tokenizer(text, return_tensors='pt')
+            input_ids = inputs.input_ids.to(dtype=torch.long, device=self.embedding_model.device)
 
             # Get the model's dtype from its parameters for the attention mask
-            model_dtype = next(self.model.parameters()).dtype
+            model_dtype = next(self.embedding_model.parameters()).dtype
 
             # Create an attention mask with matching dtype
             attention_mask = torch.zeros(
@@ -269,7 +306,7 @@ class LLMApi:
 
             # Get model outputs
             with torch.no_grad():
-                outputs = self.model(
+                outputs = self.embedding_model(
                     input_ids=input_ids,
                     attention_mask=attention_mask,
                     output_hidden_states=True,

main/app.py

@@ -1,5 +1,4 @@
 import yaml
-import sys
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from .routes import router, init_router

main/routes.py

@@ -1,3 +1,7 @@
+# routes.py for the LLM Engine.
+# This file contains the FastAPI routes for the LLM Engine API.
+# It includes routes for generating text, generating embeddings, checking system status, and validating system configuration.
+
 from fastapi import APIRouter, HTTPException
 from pydantic import BaseModel
 from typing import Optional, List, Dict, Union
@@ -51,12 +55,12 @@ class ValidationResponse(BaseModel):
 @router.get("/system/validate",
             response_model=ValidationResponse,
             summary="Validate System Configuration",
-            description="Validates system configuration, folders, and model setup")
+            description="Validates system configuration, folders, and model setup for both generation and embedding models")
 async def validate_system():
     """
     Validates:
     - Configuration parameters
-    - Model setup
+    - Model setup for both generation and embedding models
     - Folder structure
     - Required permissions
     """
@@ -88,20 +92,31 @@ async def validate_system():
     # Validate model setup
     try:
         model_status = {
-            "model_files_exist": False,
-            "model_loadable": False,
+            "generation_model_files_exist": False,
+            "generation_model_loadable": False,
+            "embedding_model_files_exist": False,
+            "embedding_model_loadable": False,
             "tokenizer_valid": False
         }
 
-        if api.model_name:
-            model_path = api.models_path / api.model_name.split('/')[-1]
-            model_status["model_files_exist"] = validate_model_path(model_path)
+        if api.generation_model_name:
+            gen_model_path = api.models_path / api.generation_model_name.split('/')[-1]
+            model_status["generation_model_files_exist"] = validate_model_path(gen_model_path)
+            model_status["generation_model_loadable"] = api.generation_model is not None
+
+        if api.embedding_model_name:
+            emb_model_path = api.models_path / api.embedding_model_name.split('/')[-1]
+            model_status["embedding_model_files_exist"] = validate_model_path(emb_model_path)
+            model_status["embedding_model_loadable"] = api.embedding_model is not None
 
-            if not model_status["model_files_exist"]:
-                issues.append("Model files are missing or incomplete")
+        model_status["tokenizer_valid"] = (
+                api.tokenizer is not None and api.embedding_tokenizer is not None
+        )
 
-            model_status["model_loadable"] = api.model is not None
-            model_status["tokenizer_valid"] = api.tokenizer is not None
+        if not model_status["generation_model_files_exist"]:
+            issues.append("Generation model files are missing or incomplete")
+        if not model_status["embedding_model_files_exist"]:
+            issues.append("Embedding model files are missing or incomplete")
 
     except Exception as e:
         logger.error(f"Model validation failed: {str(e)}")
@@ -110,9 +125,12 @@ async def validate_system():
 
     # Validate folder structure and permissions
     try:
-        folder_status = {"models_folder": api.models_path.exists(), "cache_folder": api.cache_path.exists(),
-                         "logs_folder": Path(api.base_path / "logs").exists(), "write_permissions": False}
-
+        folder_status = {
+            "models_folder": api.models_path.exists(),
+            "cache_folder": api.cache_path.exists(),
+            "logs_folder": Path(api.base_path / "logs").exists(),
+            "write_permissions": False
+        }
 
         # Test write permissions by attempting to create a test file
         test_file = api.models_path / ".test_write"
@@ -148,7 +166,6 @@ async def validate_system():
     logger.info(f"System validation completed with status: {overall_status}")
     return validation_response
 
-
 @router.get("/system/status",
             response_model=SystemStatusResponse,
             summary="Check System Status",
@@ -224,12 +241,16 @@ async def check_system():
 
     # Check Model Status
     try:
-        current_model_path = api.models_path / api.model_name.split('/')[-1] if api.model_name else None
         status.model = {
-            "is_loaded": api.model is not None,
-            "current_model": api.model_name,
-            "is_valid": validate_model_path(current_model_path) if current_model_path else False,
-            "has_chat_template": api.has_chat_template() if api.model else False
+            "generation_model": {
+                "is_loaded": api.generation_model is not None,
+                "current_model": api.generation_model_name,
+                "has_chat_template": api.has_chat_template() if api.generation_model else False
+            },
+            "embedding_model": {
+                "is_loaded": api.embedding_model is not None,
+                "current_model": api.embedding_model_name
+            }
         }
         logger.debug(f"Model status retrieved: {status.model}")
     except Exception as e:
@@ -239,7 +260,6 @@ async def check_system():
     logger.info("System status check completed")
     return status
 
-
 @router.post("/generate")
 async def generate_text(request: GenerateRequest):
     """Generate text response from prompt"""
@@ -256,7 +276,6 @@ async def generate_text(request: GenerateRequest):
         logger.error(f"Error in generate_text endpoint: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))
 
-
 @router.post("/generate/stream")
 async def generate_stream(request: GenerateRequest):
     """Generate streaming text response from prompt"""
@@ -271,7 +290,6 @@ async def generate_stream(request: GenerateRequest):
         logger.error(f"Error in generate_stream endpoint: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))
 
-
 @router.post("/embedding", response_model=EmbeddingResponse)
 async def generate_embedding(request: EmbeddingRequest):
     """Generate embedding vector from text"""
@@ -287,7 +305,6 @@ async def generate_embedding(request: EmbeddingRequest):
         logger.error(f"Error in generate_embedding endpoint: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))
 
-
 @router.post("/model/download",
              summary="Download default or specified model",
              description="Downloads model files. Uses default model from config if none specified.")
@@ -332,18 +349,30 @@ async def initialize_model(model_name: Optional[str] = None):
         logger.error(f"Error initializing model: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))
 
-
-@router.get("/models/status")
-async def get_model_status():
-    """Get current model status"""
+@router.post("/model/initialize/embedding",
+             summary="Initialize embedding model",
+             description="Initialize a separate model specifically for generating embeddings")
+async def initialize_embedding_model(model_name: Optional[str] = None):
+    """Initialize a model specifically for embeddings"""
     try:
-        status = {
-            "model_loaded": api.model is not None,
-            "current_model": api.model_name if api.model_name else None,
-            "has_chat_template": api.has_chat_template() if api.model else False
+        # Use model name from config if none provided
+        embedding_model = model_name or config["model"]["defaults"].get("embedding_model_name")
+        if not embedding_model:
+            raise HTTPException(
+                status_code=400,
+                detail="No embedding model specified and no default found in config"
+            )
+
+        logger.info(f"Received request to initialize embedding model: {embedding_model}")
+
+        api.initialize_embedding_model(embedding_model)
+        logger.info(f"Successfully initialized embedding model: {embedding_model}")
+
+        return {
+            "status": "success",
+            "message": f"Embedding model {embedding_model} initialized",
+            "model_name": embedding_model
         }
-        logger.info(f"Retrieved model status: {status}")
-        return status
     except Exception as e:
-        logger.error(f"Error getting model status: {str(e)}")
+        logger.error(f"Error initializing embedding model: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))