Update space asset usage (chroma binary)

infer_hybrid_RAG.py

@@ -10,16 +10,17 @@ from typing import Iterator
 
 sys.path.insert(0, str(Path(__file__).resolve().parent / "src"))
 
-from numen_scriptorium.inference.qwen import generate, load_model, stream_generate
-
 
 RAG_BASE_MODEL = os.getenv("NS_RAG_BASE_MODEL", os.getenv("NS_BASE_MODEL", "Qwen/Qwen2.5-7B-Instruct"))
 RAG_ADAPTER = os.getenv("NS_RAG_ADAPTER", os.getenv("NS_ADAPTER", "ICGenAIShare06/boh-qlora-adapter/best")).strip() or None
 RAG_USE_4BIT = os.getenv("NS_RAG_USE_4BIT", os.getenv("NS_USE_4BIT", "1")) == "1"
-RAG_CHROMA_DIR = os.getenv("NS_RAG_CHROMA_DIR", "chroma_data")
+RAG_LOCAL_DIR = os.getenv("NS_RAG_LOCAL_DIR", os.getenv("NS_RAG_CHROMA_DIR", "chroma_data"))
+RAG_CHROMA_DIR = RAG_LOCAL_DIR
 RAG_COLLECTION = os.getenv("NS_RAG_COLLECTION", "mansus_lore")
 RAG_ALIAS_FILE = os.getenv("NS_RAG_ALIAS_FILE", "data/hours_merged.json")
 RAG_EMBED_MODEL = os.getenv("NS_RAG_EMBED_MODEL", "moka-ai/m3e-base")
+RAG_ASSETS_REPO = os.getenv("NS_RAG_ASSETS_REPO", "ICGenAIShare06/numen-scriptorium-rag-assets")
+RAG_ASSETS_FILE = os.getenv("NS_RAG_ASSETS_FILE", "chroma.sqlite3")
 
 
 def _resolve_repo_path(path_like: str) -> Path:
@@ -29,6 +30,37 @@ def _resolve_repo_path(path_like: str) -> Path:
     return Path(__file__).resolve().parent / p
 
 
+# Runtime DB download avoids committing binary sqlite into the Space repo.
+# Control via env vars:
+# - NS_RAG_ASSETS_REPO (HF dataset repo)
+# - NS_RAG_ASSETS_FILE (remote sqlite filename)
+# - NS_RAG_LOCAL_DIR (local storage directory)
+def ensure_chroma_sqlite(local_dir: str | Path) -> Path:
+    from huggingface_hub import hf_hub_download
+
+    target_dir = Path(local_dir)
+    target_dir.mkdir(parents=True, exist_ok=True)
+    local_sqlite_path = target_dir / "chroma.sqlite3"
+
+    if local_sqlite_path.exists() and local_sqlite_path.stat().st_size > 0:
+        return local_sqlite_path
+
+    downloaded_path = hf_hub_download(
+        repo_id=RAG_ASSETS_REPO,
+        filename=RAG_ASSETS_FILE,
+        repo_type="dataset",
+    )
+    local_sqlite_path.write_bytes(Path(downloaded_path).read_bytes())
+    return local_sqlite_path
+
+
+@lru_cache(maxsize=1)
+def _get_qwen_fns():
+    from numen_scriptorium.inference.qwen import generate, load_model, stream_generate
+
+    return generate, load_model, stream_generate
+
+
 class HybridRetriever:
     def __init__(self, chroma_dir: str, collection_name: str, alias_file: str, embed_model: str):
         import chromadb
@@ -36,6 +68,7 @@ class HybridRetriever:
         from sentence_transformers import SentenceTransformer
 
         chroma_path = _resolve_repo_path(chroma_dir)
+        ensure_chroma_sqlite(chroma_path)
         self.chroma_client = chromadb.PersistentClient(path=str(chroma_path))
         self.collection = self.chroma_client.get_or_create_collection(name=collection_name)
 
@@ -92,6 +125,7 @@ def get_hybrid_retriever() -> HybridRetriever:
 
 @lru_cache(maxsize=1)
 def get_rag_runtime():
+    _, load_model, _ = _get_qwen_fns()
     return load_model(base_model=RAG_BASE_MODEL, lora_dir=RAG_ADAPTER, use_4bit=RAG_USE_4BIT)
 
 
@@ -155,6 +189,7 @@ def rag_answer(
         return ""
 
     tokenizer, model = get_rag_runtime()
+    generate, _, _ = _get_qwen_fns()
     rag_instruction = _build_rag_instruction(instruction, rag_dict, vector_context)
     return generate(
         tokenizer=tokenizer,
@@ -187,6 +222,7 @@ def rag_answer_stream(
         return
 
     tokenizer, model = get_rag_runtime()
+    _, _, stream_generate = _get_qwen_fns()
     rag_instruction = _build_rag_instruction(instruction, rag_dict, vector_context)
     yield from stream_generate(
         tokenizer=tokenizer,

requirements.txt

@@ -6,4 +6,5 @@ accelerate>=0.33.0
 sentencepiece>=0.2.0
 bitsandbytes
 chromadb>=0.5.0
-sentence-transformers>=3.0.1
+sentence-transformers>=3.0.1
+huggingface_hub>=0.24.0