Update text_embedder.py

text_embedder.py

@@ -23,13 +23,18 @@ from __future__ import annotations
 import os
 import sys
 from typing import Optional
-
+import logging
 import numpy as np
 import torch
 import torch.nn.functional as F
 from pathlib import Path
 from huggingface_hub import snapshot_download
 
+
+
+log = logging.getLogger("text_embedder")
+logging.basicConfig(level=logging.INFO, format="%(levelname)s [%(name)s] %(message)s")
+
 # ---------------------------------------------------------------------------
 # Make the local 'detree' package importable
 # ---------------------------------------------------------------------------
@@ -39,11 +44,11 @@ if _current_dir not in sys.path:
 
 try:
     from detree.model.text_embedding import TextEmbeddingModel
+    log.info("TextEmbeddingModel imported successfully.")
 except ImportError as _e:
-    print(f"Warning: could not import TextEmbeddingModel ({_e}). Text embedding will return zeros.")
+    log.error(f"Could not import TextEmbeddingModel: {_e}")
     TextEmbeddingModel = None
 
-
 # ---------------------------------------------------------------------------
 # Configuration
 # ---------------------------------------------------------------------------
@@ -56,6 +61,8 @@ REPO_ID = "MAS-AI-0000/Authentica"
 TEXT_SUBFOLDER = "Lib/Models/Text"   # where config.json/model.safetensors live in the repo
 EMBEDDING_FILE = "priori1_center10k.pt" 
 _TEXT_DIR = None
+log.info(f"[config] device={DEVICE!r}  max_length={MAX_LENGTH}  pooling={POOLING!r}")
+
 
 try:
     # download a local snapshot of just the Text folder and point _TEXT_DIR at it
@@ -78,18 +85,20 @@ except Exception as e:
 _model:     Optional[object] = None
 _tokenizer: Optional[object] = None
 
-
 def _init() -> None:
     global _model, _tokenizer
 
+    log.info("_init: starting TextEmbedder initialisation.")
+
     if TextEmbeddingModel is None:
-        print("TextEmbedder: TextEmbeddingModel unavailable — embedding disabled.")
+        log.error("_init: TextEmbeddingModel is None — check import error above. Embedding disabled.")
         return
 
     if not os.path.exists(_TEXT_DIR):
-        print(f"TextEmbedder: model directory not found at {_TEXT_DIR!r} — embedding disabled.")
+        log.error(f"_init: model directory not found at {_TEXT_DIR!r} — embedding disabled.")
         return
 
+    log.info(f"_init: loading TextEmbeddingModel from {_TEXT_DIR!r} on device={DEVICE!r} ...")
     try:
         _model = TextEmbeddingModel(
             _TEXT_DIR,
@@ -99,9 +108,10 @@ def _init() -> None:
         ).to(DEVICE)
         _model.eval()
         _tokenizer = _model.tokenizer
-        print(f"TextEmbedder: model loaded from {_TEXT_DIR!r}")
+        log.info(f"_init: model loaded OK. tokenizer type={type(_tokenizer).__name__!r}")
+        log.info(f"_init: model device={next(_model.parameters()).device}")
     except Exception as exc:
-        print(f"TextEmbedder: error loading model: {exc}")
+        log.exception(f"_init: error loading model: {exc}")
 
 
 _init()
@@ -133,23 +143,34 @@ def get_text_embedding(
         ``np.ndarray`` of shape ``(1, embedding_dim)`` and dtype float32.
     """
     if _model is None or _tokenizer is None:
+        log.error("get_text_embedding: model or tokenizer is None — returning zeros. Check _init logs.")
         return np.zeros((1, 1), dtype=np.float32)
 
-    encoded = _tokenizer(
-        [text],
-        return_tensors="pt",
-        max_length=max_length,
-        padding="max_length",
-        truncation=True,
-    )
-    encoded = {k: v.to(DEVICE) for k, v in encoded.items()}
+    log.info(f"get_text_embedding: input text length={len(text)} chars, layer={layer}")
+    try:
+        encoded = _tokenizer(
+            [text],
+            return_tensors="pt",
+            max_length=max_length,
+            padding="max_length",
+            truncation=True,
+        )
+        log.info(f"get_text_embedding: tokenised keys={list(encoded.keys())}  "
+                 f"input_ids shape={encoded['input_ids'].shape}")
+        encoded = {k: v.to(DEVICE) for k, v in encoded.items()}
 
-    # Shape returned by model with hidden_states=True: (batch, num_layers, dim)
-    embeddings = _model(encoded, hidden_states=True)
-    embeddings = F.normalize(embeddings, dim=-1)         # normalise feature dim
+        # Shape returned by model with hidden_states=True: (batch, num_layers, dim)
+        embeddings = _model(encoded, hidden_states=True)
+        log.info(f"get_text_embedding: raw embeddings shape={tuple(embeddings.shape)}")
+        embeddings = F.normalize(embeddings, dim=-1)         # normalise feature dim
 
-    # embeddings: (1, num_layers, dim)  →  select layer  →  (1, dim)
-    selected = embeddings[:, layer, :]                   # supports negative indexing
+        # embeddings: (1, num_layers, dim)  →  select layer  →  (1, dim)
+        selected = embeddings[:, layer, :]                   # supports negative indexing
+        log.info(f"get_text_embedding: selected layer={layer}  output shape={tuple(selected.shape)}  "
+                 f"norm={selected.norm(dim=-1).item():.4f}")
+    except Exception as exc:
+        log.exception(f"get_text_embedding: failed during inference: {exc}")
+        return np.zeros((1, 1), dtype=np.float32)
 
     return selected.cpu().numpy().astype(np.float32)