Update level_classifier_tool_2.py

level_classifier_tool_2.py

@@ -8,9 +8,6 @@ from transformers import AutoTokenizer, AutoModel
 #import tensorflow
 Agg = Literal["mean", "max", "topk_mean"]
 
-
-# --------------------------- Embedding backend ---------------------------
-
 @dataclass
 class HFEmbeddingBackend:
     """
@@ -18,18 +15,36 @@ class HFEmbeddingBackend:
     Uses mean pooling over last_hidden_state and L2 normalizes the result.
     """
     model_name: str = "google/embeddinggemma-300m"
-    device =  "cuda" if torch.cuda.is_available() else "cpu"
-    TOK = AutoTokenizer.from_pretrained(model_name)
-    MODEL = AutoModel.from_pretrained(model_name)
-    MODEL.to(device).eval()
-
-    def encode(self, texts: Iterable[str], batch_size: int = 32) -> "tuple[torch.Tensor, list[str]]":
+    device: str = field(default_factory=lambda: "cuda" if torch.cuda.is_available() else "cpu")
+    # Lazy-initialized in __post_init__
+    TOK: Any = field(init=False, repr=False)
+    MODEL: Any = field(init=False, repr=False)
+
+    def __post_init__(self):
+        
+        os.environ.setdefault("SPACES_ZERO_DISABLED", "1")
+
+       try:
+            torch.backends.cuda.sdp_kernel(enable_math=True, enable_flash=False, enable_mem_efficient=False)
+        except Exception:
+            pass
+
+        self.TOK = AutoTokenizer.from_pretrained(self.model_name)
+        self.MODEL = AutoModel.from_pretrained(self.model_name, attn_implementation="eager")
+        try:
+            self.MODEL.config.attn_implementation = "eager"
+        except Exception:
+            pass
+
+        self.MODEL.to(self.device).eval()
+
+    def encode(self, texts: Iterable[str], batch_size: int = 32) -> "Tuple[torch.Tensor, List[str]]":
         """
         Returns (embeddings, texts_list). Embeddings have shape [N, D] and are unit-normalized.
         """
         texts_list = list(texts)
         if not texts_list:
-            return torch.empty((0, self.MODEL.config.hidden_size)), []  # type: ignore
+            return torch.empty((0, self.MODEL.config.hidden_size)), []
 
         all_out = []
         with torch.inference_mode():
@@ -39,18 +54,18 @@ class HFEmbeddingBackend:
                 out = self.MODEL(**enc)
                 last = out.last_hidden_state  # [B, T, H]
                 mask = enc["attention_mask"].unsqueeze(-1)  # [B, T, 1]
-                # mean pool
+
+                # Mean pool
                 summed = (last * mask).sum(dim=1)
                 counts = mask.sum(dim=1).clamp(min=1)
                 pooled = summed / counts
+
                 # L2 normalize
                 pooled = pooled / pooled.norm(dim=1, keepdim=True).clamp(min=1e-12)
                 all_out.append(pooled.cpu())
-        embs = torch.cat(all_out, dim=0) if all_out else torch.empty((0, self.MODEL.config.hidden_size))  # type: ignore
-        return embs, texts_list
 
-
-# --------------------------- Utilities ---------------------------
+        embs = torch.cat(all_out, dim=0) if all_out else torch.empty((0, self.MODEL.config.hidden_size))  # type: ignore
+        return embs, texts_list----------------- Utilities ---------------------------
 
 def _normalize_whitespace(s: str) -> str:
     return " ".join(s.strip().split())