Update rag_engine.py

rag_engine.py

@@ -1,10 +1,26 @@
 from __future__ import annotations
+
 import html
-import spaces
-import re
 import os
+import re
 import threading
 from dataclasses import dataclass
+
+IS_HF_SPACE = bool(os.getenv("SPACE_ID"))
+
+if IS_HF_SPACE:
+    import spaces
+else:
+    
+    class _SpacesShim:
+        @staticmethod
+        def GPU(duration: int = 60, **_kwargs):
+            def decorator(fn):
+                return fn
+            return decorator
+
+    spaces = _SpacesShim()
+
 import faiss
 import numpy as np
 import pandas as pd
@@ -83,7 +99,7 @@ class KnowledgeBase:
                 self.dimension
             )
 
-        # Prevent accidental mixing of embeddings with different vector dimensions
+        # prevent accidental mixing of embeddings with different vector dimensions
         if embeddings.shape[1] != self.dimension:
             raise ValueError(
                 "Embedding dimension changed during this session"
@@ -210,18 +226,63 @@ _LOCAL_MODEL_LOCK = threading.Lock()
 
 def get_embedding_device() -> str:
     """
-     embedding device
+    Return the device used by the embedding model.
+
+    On a Hugging Face ZeroGPU Space the model is registered on CUDA at
+    module scope. ZeroGPU intercepts that placement and attaches a real
+    GPU only while a @spaces.GPU function is running.
+
+    During local development, use the local NVIDIA GPU when available
+    and otherwise fall back to CPU.
     """
 
-    if os.getenv("SPACES_ZERO_GPU"):
-        return "cpu"
+    if IS_HF_SPACE:
+        return "cuda"
 
     return "cuda" if torch.cuda.is_available() else "cpu"
 
 
+def _load_embedder() -> SentenceTransformer:
+    """
+    Construct the embedding model.
+
+    ZeroGPU requires CUDA model placement at module scope rather than
+    lazy-loading the model inside a GPU-decorated request.
+    """
+
+    device = get_embedding_device()
+
+    print(
+        f"[embedding] Loading {EMBEDDING_MODEL_ID} on {device}",
+        flush=True,
+    )
+
+    model = SentenceTransformer(
+        EMBEDDING_MODEL_ID,
+        device=device,
+    )
+
+    print(
+        "[embedding] Model ready.",
+        flush=True,
+    )
+
+    return model
+
+
+if IS_HF_SPACE:
+    _EMBEDDER = _load_embedder()
+
+
 def get_embedder() -> SentenceTransformer:
     """
-    Load the embedding model once and reuse it
+    Return the shared embedding model.
+
+    Hugging Face Space:
+        The model was loaded at module scope for ZeroGPU.
+
+    Local machine:
+        Load once on first use and reuse it afterward.
     """
 
     global _EMBEDDER
@@ -230,36 +291,22 @@ def get_embedder() -> SentenceTransformer:
         return _EMBEDDER
 
     with _EMBEDDER_LOCK:
-
         if _EMBEDDER is None:
-
-            device = get_embedding_device()
-
-            print(
-                f"[embedding] Loading "
-                f"{EMBEDDING_MODEL_ID} "
-                f"on {device}",
-                flush=True,
-            )
-
-            _EMBEDDER = SentenceTransformer(
-                EMBEDDING_MODEL_ID,
-                device=device,
-            )
-
-            print(
-                "[embedding] Model ready.",
-                flush=True,
-            )
+            _EMBEDDER = _load_embedder()
 
     return _EMBEDDER
 
 
 def _encode(
     texts: list[str],
+    *,
+    batch_size: int = 16,
 ) -> np.ndarray:
     """
-    Encode text passages into normalized embeddings
+    Encode text into normalized float32 vectors for FAISS.
+
+    This helper performs the actual model computation. On ZeroGPU it
+    must only be reached from a @spaces.GPU-decorated function.
     """
 
     if not texts:
@@ -271,40 +318,46 @@ def _encode(
 
     vectors = model.encode(
         texts,
-        batch_size=8,
+        batch_size=int(batch_size),
         normalize_embeddings=True,
         convert_to_numpy=True,
         show_progress_bar=False,
     )
 
-    # FAISS expects float32 arrays
     vectors = vectors.astype("float32")
 
+    # return ordinary CPU/NumPy data across the ZeroGPU process boundary.
     return np.ascontiguousarray(vectors)
 
 
+@spaces.GPU(duration=120)
 def embed_passages(
     texts: list[str],
 ) -> np.ndarray:
     """
-    Embed document passages
+    Embed document/FAR passages.
+
+    Only the embedding computation runs on ZeroGPU. The FAISS index and
+    session knowledge base remain in the main application process.
     """
 
-    return _encode(texts)
+    return _encode(
+        texts,
+        batch_size=16,
+    )
 
 
+@spaces.GPU(duration=30)
 def embed_query(
     query: str,
 ) -> np.ndarray:
     """
-    Embed a search query
+    Embed a retrieval query on ZeroGPU.
 
-    Qwen embedding models can expose a query-specific prompt
-    If unavailable, a generic retrieval instruction is used
+    Qwen embedding models may expose a query-specific prompt. If that
+    prompt is unavailable, use a generic retrieval instruction.
     """
 
-    model = get_embedder()
-
     clean = query.strip()
 
     if not clean:
@@ -312,6 +365,8 @@ def embed_query(
             "Query cannot be empty"
         )
 
+    model = get_embedder()
+
     prompts = getattr(
         model,
         "prompts",
@@ -319,7 +374,6 @@ def embed_query(
     )
 
     if "query" in prompts:
-
         vectors = model.encode(
             [clean],
             prompt_name="query",
@@ -337,14 +391,14 @@ def embed_query(
             vectors
         )
 
-    # compatibility path for models
     instructed = (
         "Represent this sentence for searching "
         f"relevant passages: {clean}"
     )
 
     return _encode(
-        [instructed]
+        [instructed],
+        batch_size=1,
     )
 
 
@@ -582,7 +636,7 @@ def build_context(
             + chunk.text.strip()
         )
 
-        # Stop adding context once the configured character budget has been reached
+        # stop adding context once the configured character budget has been reached
         if (
             used_chars + len(block) > max_chars
             and blocks
@@ -607,6 +661,12 @@ def _get_local_model():
     (CUDA is used automatically when available)
     """
 
+    if IS_HF_SPACE:
+        raise RuntimeError(
+            "Local LLM fallback is disabled on this ZeroGPU Space. "
+            "Use LLM_BACKEND=hf_api with HF_TOKEN configured as a Space Secret."
+        )
+
     global _LOCAL_MODEL
     global _LOCAL_TOKENIZER
 
@@ -851,17 +911,18 @@ def _generate_local(
     return answer.strip()
 
 
-# Backend selection
+# backend selection
 
 def choose_backend() -> str:
     """
-    Choose between hosted Hugging Face inference
-    and local generation
+    Choose the generation backend.
+
+    Hugging Face ZeroGPU deployment:
+        Use hosted Hugging Face inference for generation. ZeroGPU is
+        reserved for the embedding model.
 
-    LLM_BACKEND values:
-        hf_api
-        local
-        auto
+    Local development:
+        Respect LLM_BACKEND = hf_api, local, or auto.
     """
 
     configured = (
@@ -870,22 +931,29 @@ def choose_backend() -> str:
         .lower()
     )
 
+    if IS_HF_SPACE:
+        if not HF_TOKEN:
+            raise RuntimeError(
+                "HF_TOKEN is required on the Hugging Face Space. "
+                "Configure it as a Space Secret and use "
+                "LLM_BACKEND=hf_api."
+            )
+
+        return "hf_api"
+
     if configured in {
         "hf_api",
         "local",
     }:
         return configured
 
-    # Auto behavior:
-    # use HF API when a token exists,
-    # otherwise use the local model
     if HF_TOKEN:
         return "hf_api"
 
     return "local"
 
 
-# Citation validation
+# citation validation
 
 def _extract_citation_ids(
     answer: str,
@@ -902,7 +970,7 @@ def _extract_citation_ids(
     )
 
 
-# Main RAG answer function
+# main RAG answer function
 
 def answer_question(
     question: str,
@@ -1206,4 +1274,4 @@ def render_sources(
         "<div class='sources-grid'>"
         + "".join(cards)
         + "</div>"
-    )
+    )