Update level_classifier_tool.py

level_classifier_tool.py

@@ -22,7 +22,9 @@ class HFEmbeddingBackend:
     Uses mean pooling over last_hidden_state and L2 normalizes the result.
     """
     model_name: str = "sentence-transformers/all-MiniLM-L6-v2"
-    device: Optional[str] = None  # "cuda" | "cpu" | None -> auto
+    # "cuda" | "cpu" | None -> (env or "cpu")
+    # We default to CPU on Spaces to avoid ZeroGPU device mixups.
+    device: Optional[str] = None
 
     def _lazy_import(self) -> None:
         global _TOK, _MODEL, _TORCH
@@ -33,24 +35,27 @@ class HFEmbeddingBackend:
             from transformers import AutoTokenizer, AutoModel  # type: ignore
             _TOK = AutoTokenizer.from_pretrained(self.model_name)
             _MODEL = AutoModel.from_pretrained(self.model_name)
-        dev = self.device or ("cuda" if _TORCH.cuda.is_available() else "cpu")
+        # Prefer explicit device -> env override -> default to CPU
+        dev = self.device or os.getenv("EMBEDDING_DEVICE") or "cpu"
         _MODEL.to(dev).eval()
         self.device = dev
 
     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.
+        Returns (embeddings, texts_list). Embeddings are a CPU torch.Tensor [N, D], unit-normalized.
         """
         self._lazy_import()
         torch = _TORCH  # local alias
         texts_list = list(texts)
         if not texts_list:
+            # Hidden size available after _lazy_import
             return torch.empty((0, _MODEL.config.hidden_size)), []  # type: ignore
 
         all_out = []
         with torch.inference_mode():
             for i in range(0, len(texts_list), batch_size):
                 batch = texts_list[i:i + batch_size]
+                # Tokenize and move to model device
                 enc = _TOK(batch, padding=True, truncation=True, return_tensors="pt").to(self.device)  # type: ignore
                 out = _MODEL(**enc)
                 last = out.last_hidden_state  # [B, T, H]
@@ -61,7 +66,9 @@ class HFEmbeddingBackend:
                 pooled = summed / counts
                 # L2 normalize
                 pooled = pooled / pooled.norm(dim=1, keepdim=True).clamp(min=1e-12)
+                # Collect on CPU for downstream ops
                 all_out.append(pooled.cpu())
+
         embs = torch.cat(all_out, dim=0) if all_out else torch.empty((0, _MODEL.config.hidden_size))  # type: ignore
         return embs, texts_list
 
@@ -102,16 +109,22 @@ def build_phrase_index(
         cur += len(plist)
         spans.append((lvl, start, cur))
 
-    embs, _ = backend.encode(all_texts)
+    embs, _ = backend.encode(all_texts)  # embs is a CPU torch.Tensor [N, D]
+
     # Slice embeddings back into level buckets
     torch = _TORCH
     embeddings_by_level: Dict[str, "Any"] = {}
     for lvl, start, end in spans:
-        embeddings_by_level[lvl] = embs[start:end] if end > start else torch.empty((0, embs.shape[1]))  # type: ignore
+        if end > start:
+            embeddings_by_level[lvl] = embs[start:end]  # torch.Tensor slice [n_i, D]
+        else:
+            embeddings_by_level[lvl] = torch.empty((0, embs.shape[1]))  # type: ignore
 
-    return PhraseIndex(phrases_by_level={lvl: list(pl) for lvl, pl in cleaned.items()},
-                       embeddings_by_level=embeddings_by_level,
-                       model_name=backend.model_name)
+    return PhraseIndex(
+        phrases_by_level={lvl: list(pl) for lvl, pl in cleaned.items()},
+        embeddings_by_level=embeddings_by_level,
+        model_name=backend.model_name
+    )
 
 
 def _aggregate_sims(
@@ -153,64 +166,20 @@ def classify_levels_phrases(
     """
     Score a question against Bloom's taxonomy and DOK (Depth of Knowledge)
     using cosine similarity to level-specific anchor phrases.
-
-    Parameters
-    ----------
-    question : str
-        The input question or prompt.
-    blooms_phrases : dict[str, Iterable[str]]
-        Mapping level -> list of anchor phrases for Bloom's.
-    dok_phrases : dict[str, Iterable[str]]
-        Mapping level -> list of anchor phrases for DOK.
-    model_name : str
-        Hugging Face model name for text embeddings. Ignored when `backend` provided.
-    agg : {"mean","max","topk_mean"}
-        Aggregation over phrase similarities within a level.
-    topk : int
-        Used only when `agg="topk_mean"`.
-    preprocess : Optional[Callable[[str], str]]
-        Preprocessing function for the question string. Defaults to whitespace normalization.
-    backend : Optional[HFEmbeddingBackend]
-        Injected embedding backend. If not given, one is constructed.
-    prebuilt_bloom_index, prebuilt_dok_index : Optional[PhraseIndex]
-        If provided, reuse precomputed phrase embeddings to avoid re-encoding.
-    return_phrase_matches : bool
-        If True, returns per-level top contributing phrases.
-
-    Returns
-    -------
-    dict
-        {
-          "question": ...,
-          "model_name": ...,
-          "blooms": {
-              "scores": {level: float, ...},
-              "best_level": str,
-              "best_score": float,
-              "top_phrases": {level: [(phrase, sim_float), ...], ...}  # only if return_phrase_matches
-          },
-          "dok": {
-              "scores": {level: float, ...},
-              "best_level": str,
-              "best_score": float,
-              "top_phrases": {level: [(phrase, sim_float), ...], ...}  # only if return_phrase_matches
-          },
-          "config": {"agg": agg, "topk": topk if agg=='topk_mean' else None}
-        }
     """
     preprocess = preprocess or _default_preprocess
     question_clean = preprocess(question)
 
-    # Prepare backend
+    # Prepare backend (defaults to CPU)
     be = backend or HFEmbeddingBackend(model_name=model_name)
 
     # Build / reuse indices
     bloom_index = prebuilt_bloom_index or build_phrase_index(be, blooms_phrases)
     dok_index = prebuilt_dok_index or build_phrase_index(be, dok_phrases)
 
-    # Encode question
+    # Encode question -> CPU torch.Tensor [1, D]
     q_emb, _ = be.encode([question_clean])
-    q_emb = q_emb[0:1]  # [1, D]
+    q_emb = q_emb[0:1]
     torch = _TORCH
 
     def _score_block(index: PhraseIndex) -> Tuple[Dict[str, float], Dict[str, List[Tuple[str, float]]]]:
@@ -218,12 +187,13 @@ def classify_levels_phrases(
         top_contribs: Dict[str, List[Tuple[str, float]]] = {}
 
         for lvl, phrases in index.phrases_by_level.items():
-            embs = index.embeddings_by_level[lvl]  # [N, D]
+            embs = index.embeddings_by_level[lvl]  # torch.Tensor [N, D]
             if embs.numel() == 0:
                 scores[lvl] = float("nan")
                 top_contribs[lvl] = []
                 continue
-            sims = (q_emb @ embs.T).squeeze(0)  # cosine sim due to L2 norm
+            # cosine similarity since embs and q_emb are unit-normalized
+            sims = (q_emb @ embs.T).squeeze(0)
             scores[lvl] = _aggregate_sims(sims, agg, topk)
             if return_phrase_matches:
                 k = min(5, sims.numel())