Upload folder using huggingface_hub

rag_system/guardrails.py

@@ -277,16 +277,29 @@ def _check_query_llama_guard(query: str) -> GuardrailResult:
             tokenized = _llama_guard_tokenizer(prompt, return_tensors="pt")
 
         if isinstance(tokenized, dict):
-            input_ids = tokenized["input_ids"].to(_llama_guard_model.device)
+            input_ids = tokenized.get("input_ids")
             attention_mask = tokenized.get("attention_mask")
-            if attention_mask is not None:
-                attention_mask = attention_mask.to(_llama_guard_model.device)
         else:
-            input_ids = tokenized.to(_llama_guard_model.device)
+            input_ids = tokenized
             attention_mask = None
 
+        if not hasattr(input_ids, "shape"):
+            import torch
+            input_ids = torch.as_tensor(input_ids)
+        if len(getattr(input_ids, "shape", [])) == 1:
+            input_ids = input_ids.unsqueeze(0)
+        input_ids = input_ids.to(_llama_guard_model.device)
+
         if attention_mask is None:
-            attention_mask = input_ids.new_ones(input_ids.shape)
+            import torch
+            attention_mask = torch.ones_like(input_ids)
+        else:
+            if not hasattr(attention_mask, "shape"):
+                import torch
+                attention_mask = torch.as_tensor(attention_mask)
+            if len(getattr(attention_mask, "shape", [])) == 1:
+                attention_mask = attention_mask.unsqueeze(0)
+            attention_mask = attention_mask.to(_llama_guard_model.device)
 
         prompt_len = input_ids.shape[1]
         output = _llama_guard_model.generate(

rag_system/query_engine.py

@@ -72,11 +72,6 @@ def _should_preserve_exact_reference(query: str) -> bool:
     return bool(_SECTION_REF_RE.search(query) and _SECTION_HINT_RE.search(query))
 
 
-def _is_try_docs_scope(collections: list[str]) -> bool:
-    prefix = settings.try_docs_prefix
-    return bool(collections) and all(c.startswith(prefix) for c in collections)
-
-
 def _cache_collection_key(collections: list[str]) -> str:
     raw = "|".join(sorted(collections))
     return hashlib.sha1(raw.encode()).hexdigest()[:16]
@@ -254,7 +249,7 @@ async def query(
     collections = request.doc_collections or [request.collection_name]
     embedding_mode = resolve_embedding_mode_for_collections(collections, request.embedding_mode)
     mode_val = request.retrieval_mode.value if hasattr(request.retrieval_mode, "value") else str(request.retrieval_mode)
-    cache_allowed = settings.cache_enabled and _is_try_docs_scope(collections)
+    cache_allowed = settings.cache_enabled
     cache_collection_key = _cache_collection_key(collections)
     cache_params_key = _cache_params_key_v2(
         mode_val,
@@ -455,7 +450,7 @@ async def pipeline_stream_query(request: QueryRequest) -> AsyncIterator[str]:
     mode_val = request.retrieval_mode.value if hasattr(request.retrieval_mode, "value") else str(request.retrieval_mode)
     collections = request.doc_collections or [request.collection_name]
     embedding_mode = resolve_embedding_mode_for_collections(collections, request.embedding_mode)
-    cache_allowed = settings.cache_enabled and _is_try_docs_scope(collections)
+    cache_allowed = settings.cache_enabled
     cache_collection_key = _cache_collection_key(collections)
     cache_params_key = _cache_params_key_v2(
         mode_val,