Spaces:
Running
Running
| """ | |
| Background initialization system for preloading heavy components during app startup. | |
| This module handles the eager loading of embedding models, retrievers, and chunks | |
| to improve first-question response time. | |
| """ | |
| import threading | |
| import time | |
| from typing import Optional, Callable | |
| from .config import logger | |
| class BackgroundInitializer: | |
| """Manages background initialization of heavy components""" | |
| def __init__(self): | |
| self._initialization_thread: Optional[threading.Thread] = None | |
| self._initialization_complete = threading.Event() | |
| self._initialization_error: Optional[Exception] = None | |
| self._progress_callback: Optional[Callable[[str, int], None]] = None | |
| self._status = "Not started" | |
| def set_progress_callback(self, callback: Callable[[str, int], None]): | |
| """Set a callback function to receive progress updates""" | |
| self._progress_callback = callback | |
| def _update_progress(self, message: str, percentage: int): | |
| """Update progress and call callback if set""" | |
| self._status = message | |
| logger.info(f"🔄 Background Init: {message} ({percentage}%)") | |
| if self._progress_callback: | |
| try: | |
| self._progress_callback(message, percentage) | |
| except Exception as e: | |
| logger.error(f"Progress callback error: {e}") | |
| def _initialize_components(self): | |
| """Initialize all heavy components in background thread""" | |
| try: | |
| self._update_progress("Starting background initialization...", 0) | |
| # Step 1: Load embedding model (this is the heaviest component) | |
| self._update_progress("Loading embedding model...", 10) | |
| from .config import get_embedding_model | |
| embedding_model = get_embedding_model() | |
| self._update_progress("Embedding model loaded successfully", 40) | |
| # Step 2: Initialize retrievers (this will load chunks and create vector store) | |
| self._update_progress("Initializing retrievers and loading chunks...", 50) | |
| from .retrievers import _ensure_initialized | |
| _ensure_initialized() | |
| self._update_progress("Retrievers initialized successfully", 90) | |
| # Step 3: Learn medical terminology from corpus | |
| self._update_progress("Learning medical terminology from corpus...", 92) | |
| try: | |
| from .medical_terminology import learn_from_corpus | |
| from . import utils | |
| # Load chunks to learn from | |
| chunks = utils.load_chunks() | |
| if chunks: | |
| # Convert to format expected by learner | |
| documents = [{'content': chunk.page_content} for chunk in chunks[:1000]] # Limit for performance | |
| learn_from_corpus(documents) | |
| logger.info(f"Learned medical terminology from {len(documents)} documents") | |
| except Exception as e: | |
| logger.warning(f"Could not learn terminology from corpus: {e}") | |
| # Step 4: Warm up LLM (optional, lightweight) | |
| self._update_progress("Warming up LLM...", 97) | |
| from .config import get_llm | |
| llm = get_llm() | |
| self._update_progress("All components initialized successfully", 100) | |
| logger.info("✅ Background initialization completed successfully") | |
| except Exception as e: | |
| self._initialization_error = e | |
| logger.error(f"❌ Background initialization failed: {e}") | |
| self._update_progress(f"Initialization failed: {str(e)}", -1) | |
| finally: | |
| self._initialization_complete.set() | |
| def start_background_initialization(self): | |
| """Start background initialization in a separate thread""" | |
| if self._initialization_thread is not None: | |
| logger.warning("Background initialization already started") | |
| return | |
| logger.info("🚀 Starting background initialization...") | |
| self._initialization_thread = threading.Thread( | |
| target=self._initialize_components, | |
| name="BackgroundInitializer", | |
| daemon=True | |
| ) | |
| self._initialization_thread.start() | |
| def is_complete(self) -> bool: | |
| """Check if initialization is complete""" | |
| return self._initialization_complete.is_set() | |
| def wait_for_completion(self, timeout: Optional[float] = None) -> bool: | |
| """Wait for initialization to complete""" | |
| return self._initialization_complete.wait(timeout) | |
| def get_status(self) -> str: | |
| """Get current initialization status""" | |
| return self._status | |
| def get_error(self) -> Optional[Exception]: | |
| """Get initialization error if any""" | |
| return self._initialization_error | |
| def is_successful(self) -> bool: | |
| """Check if initialization completed successfully""" | |
| return self.is_complete() and self._initialization_error is None | |
| # Global initializer instance | |
| _background_initializer = BackgroundInitializer() | |
| def start_background_initialization(progress_callback: Optional[Callable[[str, int], None]] = None): | |
| """Start background initialization with optional progress callback""" | |
| if progress_callback: | |
| _background_initializer.set_progress_callback(progress_callback) | |
| _background_initializer.start_background_initialization() | |
| def wait_for_initialization(timeout: Optional[float] = None) -> bool: | |
| """Wait for background initialization to complete""" | |
| return _background_initializer.wait_for_completion(timeout) | |
| def is_initialization_complete() -> bool: | |
| """Check if background initialization is complete""" | |
| return _background_initializer.is_complete() | |
| def get_initialization_status() -> str: | |
| """Get current initialization status""" | |
| return _background_initializer.get_status() | |
| def get_initialization_error() -> Optional[Exception]: | |
| """Get initialization error if any""" | |
| return _background_initializer.get_error() | |
| def is_initialization_successful() -> bool: | |
| """Check if initialization completed successfully""" | |
| return _background_initializer.is_successful() | |