.

Author: edwin-onyx

backend/model_server/custom_models.py

@@ -1,12 +1,11 @@
 from typing import cast
-from typing import Optional
-from typing import TYPE_CHECKING
 
 import numpy as np
 import torch
 import torch.nn.functional as F
 from fastapi import APIRouter
 from huggingface_hub import snapshot_download  # type: ignore
+from setfit import SetFitModel  # type: ignore[import]
 from transformers import AutoTokenizer  # type: ignore
 from transformers import BatchEncoding  # type: ignore
 from transformers import PreTrainedTokenizer  # type: ignore
@@ -38,9 +37,6 @@
 from shared_configs.model_server_models import IntentRequest
 from shared_configs.model_server_models import IntentResponse
 
-if TYPE_CHECKING:
-    from setfit import SetFitModel  # type: ignore[import-untyped]
-
 logger = setup_logger()
 
 router = APIRouter(prefix="/custom")
@@ -51,8 +47,7 @@
 _INTENT_TOKENIZER: PreTrainedTokenizer | None = None
 _INTENT_MODEL: HybridClassifier | None = None
 
-
-_INFORMATION_CONTENT_MODEL: Optional["SetFitModel"] = None
+_INFORMATION_CONTENT_MODEL: SetFitModel | None = None
 
 _INFORMATION_CONTENT_MODEL_PROMPT_PREFIX: str = ""  # spec to model version!
 
@@ -146,10 +141,8 @@ def get_local_intent_model(
 def get_local_information_content_model(
     model_name_or_path: str = INFORMATION_CONTENT_MODEL_VERSION,
     tag: str | None = INFORMATION_CONTENT_MODEL_TAG,
-) -> "SetFitModel":
+) -> SetFitModel:
     global _INFORMATION_CONTENT_MODEL
-    from setfit import SetFitModel  # type: ignore
-
     if _INFORMATION_CONTENT_MODEL is None:
         try:
             # Calculate where the cache should be, then load from local if available

backend/model_server/encoders.py

@@ -7,8 +7,8 @@
 from fastapi import HTTPException
 from fastapi import Request
 from litellm.exceptions import RateLimitError
-from sentence_transformers import CrossEncoder
-from sentence_transformers import SentenceTransformer
+from sentence_transformers import CrossEncoder  # type: ignore
+from sentence_transformers import SentenceTransformer  # type: ignore
 
 from model_server.utils import simple_log_function_time
 from onyx.utils.logger import setup_logger
@@ -25,8 +25,8 @@
 router = APIRouter(prefix="/encoder")
 
 
-_GLOBAL_MODELS_DICT: dict[str, SentenceTransformer] = {}
-_RERANK_MODEL: Optional[CrossEncoder] = None
+_GLOBAL_MODELS_DICT: dict[str, "SentenceTransformer"] = {}
+_RERANK_MODEL: Optional["CrossEncoder"] = None
 
 # If we are not only indexing, dont want retry very long
 _RETRY_DELAY = 10 if INDEXING_ONLY else 0.1
@@ -36,13 +36,14 @@
 def get_embedding_model(
     model_name: str,
     max_context_length: int,
-) -> SentenceTransformer:
+) -> "SentenceTransformer":
     """
     Loads or returns a cached SentenceTransformer, sets max_seq_length, pins device,
     pre-warms rotary caches once, and wraps encode() with a lock to avoid cache races.
     """
+    from sentence_transformers import SentenceTransformer  # type: ignore
 
-    def _prewarm_rope(st_model: SentenceTransformer, target_len: int) -> None:
+    def _prewarm_rope(st_model: "SentenceTransformer", target_len: int) -> None:
         """
         Build RoPE cos/sin caches once on the final device/dtype so later forwards only read.
         Works by calling the underlying HF model directly with dummy IDs/attention.
@@ -101,7 +102,7 @@ def get_local_reranking_model(
 
 
 def _concurrent_embedding(
-    texts: list[str], model: SentenceTransformer, normalize_embeddings: bool
+    texts: list[str], model: "SentenceTransformer", normalize_embeddings: bool
 ) -> Any:
     """Synchronous wrapper for concurrent_embedding to use with run_in_executor."""
     for _ in range(ENCODING_RETRIES):

backend/onyx/file_processing/extract_file_text.py

@@ -15,14 +15,12 @@
 from typing import Any
 from typing import IO
 from typing import NamedTuple
+from typing import Optional
+from typing import TYPE_CHECKING
 from zipfile import BadZipFile
 
 import chardet
 import openpyxl
-from markitdown import FileConversionException
-from markitdown import MarkItDown
-from markitdown import StreamInfo
-from markitdown import UnsupportedFormatException
 from PIL import Image
 from pypdf import PdfReader
 from pypdf.errors import PdfStreamError
@@ -37,6 +35,11 @@
 from onyx.utils.file_types import WORD_PROCESSING_MIME_TYPE
 from onyx.utils.logger import setup_logger
 
+
+if TYPE_CHECKING:
+    from markitdown import MarkItDown
+
+
 logger = setup_logger()
 
 # NOTE(rkuo): Unify this with upload_files_for_chat and file_valiation.py
@@ -85,17 +88,19 @@
     "image/webp",
 ]
 
-_MARKITDOWN_CONVERTER: MarkItDown | None = None
+_MARKITDOWN_CONVERTER: Optional["MarkItDown"] = None
 
 KNOWN_OPENPYXL_BUGS = [
     "Value must be either numerical or a string containing a wildcard",
     "File contains no valid workbook part",
 ]
 
 
-def get_markitdown_converter() -> MarkItDown:
+def get_markitdown_converter() -> "MarkItDown":
     global _MARKITDOWN_CONVERTER
     if _MARKITDOWN_CONVERTER is None:
+        from markitdown import MarkItDown
+
         _MARKITDOWN_CONVERTER = MarkItDown(enable_plugins=False)
     return _MARKITDOWN_CONVERTER
 
@@ -357,6 +362,12 @@ def docx_to_text_and_images(
     of avoiding materializing the list of images in memory.
     The images list returned is empty in this case.
     """
+    from markitdown import (
+        FileConversionException,
+        StreamInfo,
+        UnsupportedFormatException,
+    )
+
     md = get_markitdown_converter()
     try:
         doc = md.convert(
@@ -393,6 +404,12 @@ def docx_to_text_and_images(
 
 
 def pptx_to_text(file: IO[Any], file_name: str = "") -> str:
+    from markitdown import (
+        FileConversionException,
+        StreamInfo,
+        UnsupportedFormatException,
+    )
+
     md = get_markitdown_converter()
     stream_info = StreamInfo(
         mimetype=PRESENTATION_MIME_TYPE, filename=file_name or None, extension=".pptx"

backend/scripts/check_lazy_imports.py

@@ -15,7 +15,7 @@
 
 logger = logging.getLogger(__name__)
 
-_MODULES_TO_LAZY_IMPORT = {"vertexai"}
+_MODULES_TO_LAZY_IMPORT = {"vertexai", "markitdown"}
 
 
 @dataclass