Upload app.py

app.py

@@ -1,9 +1,11 @@
-import gradio as gr
-from llama_cpp import Llama
-from huggingface_hub import hf_hub_download
 import os
 import pickle
 import logging
+import platform
+
+import gradio as gr
+from llama_cpp import Llama
+from huggingface_hub import hf_hub_download
 from langchain_huggingface import HuggingFaceEmbeddings
 
 # ====================== LOGGING ======================
@@ -15,16 +17,39 @@ repo_id = "robertolofaro/articles-model"
 
 BACKENDS = {
     "FAISS - RAG (HNSW)": "FAISS",
-    "Qdrant - RAG": "Qdrant"
+    "Qdrant - RAG": "Qdrant",
 }
 
-# Resolve paths relative to this file so they work in HF Spaces
-_HERE = os.path.dirname(os.path.abspath(__file__))
-METADATA_PATH  = os.path.join(_HERE, "metadata.pkl")
-FAISS_PATH     = os.path.join(_HERE, "faiss_hnsw")
-QDRANT_PATH    = os.path.join(_HERE, "qdrant_db")
+_HERE             = os.path.dirname(os.path.abspath(__file__))
+METADATA_PATH     = os.path.join(_HERE, "metadata.pkl")
+FAISS_PATH        = os.path.join(_HERE, "faiss_hnsw")
+QDRANT_PATH       = os.path.join(_HERE, "qdrant_db")
 QDRANT_COLLECTION = "articles"
 
+# ====================== GPU / HARDWARE DETECTION ======================
+# Override everything with N_GPU_LAYERS env var when you need fine control.
+# Otherwise: CUDA → all layers on GPU (-1); Apple Silicon → Metal (-1); else CPU (0).
+def _detect_gpu_layers() -> int:
+    override = os.environ.get("N_GPU_LAYERS")
+    if override is not None:
+        val = int(override)
+        logger.info("N_GPU_LAYERS override: %d", val)
+        return val
+    try:
+        import torch
+        if torch.cuda.is_available():
+            logger.info("CUDA detected — offloading all layers to GPU")
+            return -1
+    except ImportError:
+        pass
+    if platform.system() == "Darwin" and platform.machine() == "arm64":
+        logger.info("Apple Silicon / Metal detected — offloading all layers to GPU")
+        return -1
+    logger.info("No GPU detected — running on CPU only")
+    return 0
+
+N_GPU_LAYERS = _detect_gpu_layers()
+
 # ====================== LOAD METADATA ======================
 def _load_metadata():
     """Load the DataFrame from metadata.pkl; return None on any failure."""
@@ -68,21 +93,39 @@ def load_articles_for_category(category: str):
 CATEGORY_LIST = load_category_list()
 
 # ====================== LOAD LLM ======================
-model_path = hf_hub_download(
-    repo_id=repo_id,
-    filename="articles-Q4_K_M.gguf",
-    repo_type="model",
-    token=os.environ.get("HF_TOKEN"),
-)
-
-llm = Llama(
-    model_path=model_path,
-    n_ctx=8192,
-    n_threads=2,
-    n_batch=512,
-    n_ubatch=512,
-    verbose=False,
-)
+# LOCAL_MODEL_PATH env var lets you point to a local GGUF and skip the HF download.
+# N_THREADS env var overrides thread count (default: 4 on CPU, 2 on GPU).
+def _load_llm() -> Llama:
+    local_model = os.environ.get("LOCAL_MODEL_PATH")
+    if local_model and os.path.isfile(local_model):
+        model_path = local_model
+        logger.info("Using local model at %s", model_path)
+    else:
+        if local_model:
+            logger.warning("LOCAL_MODEL_PATH set but file not found (%s) — downloading from HF", local_model)
+        logger.info("Downloading model from HF hub (%s)…", repo_id)
+        model_path = hf_hub_download(
+            repo_id=repo_id,
+            filename="articles-Q4_K_M.gguf",
+            repo_type="model",
+            token=os.environ.get("HF_TOKEN"),
+        )
+
+    default_threads = 2 if N_GPU_LAYERS != 0 else 4
+    n_threads = int(os.environ.get("N_THREADS", default_threads))
+    logger.info("Llama init: n_gpu_layers=%d, n_threads=%d", N_GPU_LAYERS, n_threads)
+
+    return Llama(
+        model_path=model_path,
+        n_ctx=8192,
+        n_threads=n_threads,
+        n_batch=512,
+        n_ubatch=512,
+        n_gpu_layers=N_GPU_LAYERS,
+        verbose=False,
+    )
+
+llm = _load_llm()
 
 # ====================== RAG CACHE ======================
 vectorstores: dict = {}
@@ -113,9 +156,68 @@ def get_vectorstore(backend_name: str):
         return None
 
 
+def _rag_search(vs, query: str, k: int, article_filter: str, category_filter: str):
+    """
+    Similarity search with optional metadata filtering.
+
+    FAISS does not support dict-based server-side metadata filtering reliably
+    across langchain versions: it either ignores the filter silently or raises.
+    We therefore fetch a generous candidate pool and post-filter in Python.
+
+    Qdrant supports native dict filtering, so we pass it directly.
+    """
+    want_title    = None if article_filter  in (None, "", "All articles in category") else article_filter
+    want_category = None if category_filter in (None, "", "All categories")           else category_filter
+
+    backend_type = type(vs).__name__   # "FAISS" or "Qdrant"
+
+    if backend_type == "FAISS":
+        # Fetch a large pool, then filter in Python.
+        pool_size = min(k * 10, 80)
+        pool = vs.similarity_search(query, k=pool_size)
+
+        filtered = []
+        for doc in pool:
+            meta = doc.metadata
+            if want_title and meta.get("article_title") != want_title:
+                continue
+            if want_category and meta.get("article_category") != want_category:
+                continue
+            filtered.append(doc)
+            if len(filtered) >= k:
+                break
+
+        if not filtered and (want_title or want_category):
+            # Nothing matched the filter — warn and fall back to unfiltered results.
+            logger.warning(
+                "FAISS post-filter (title=%r, cat=%r) matched 0 docs — "
+                "returning unfiltered top-%d",
+                want_title, want_category, k,
+            )
+            return pool[:k]
+
+        logger.info(
+            "FAISS post-filter (title=%r, cat=%r) → %d/%d docs kept",
+            want_title, want_category, len(filtered), len(pool),
+        )
+        return filtered
+
+    else:
+        # Qdrant: use its native metadata filter dict.
+        filter_dict = None
+        if want_title:
+            filter_dict = {"article_title": want_title}
+        elif want_category:
+            filter_dict = {"article_category": want_category}
+
+        docs = vs.similarity_search(query, k=k, filter=filter_dict)
+        logger.info(
+            "Qdrant search (filter=%r) → %d docs", filter_dict, len(docs)
+        )
+        return docs
+
+
 # ====================== SYSTEM PROMPT ======================
-# The explicit declaration that context is injected inline prevents the model from
-# reasoning "I have no access to the vector store" during its <think> block.
 SYSTEM_PROMPT = """You are the reference expert for the articles contained in the training \
 of this model, all extracted from the website robertolofaro.com, and all focused on change.
 
@@ -143,12 +245,11 @@ def generate_response(
     max_tokens, temperature, top_p, repeat_penalty,
     suppress_thinking,
 ):
-    # Strip /nothink from the user-visible message if they typed it
+    # Strip any /nothink the user may have typed manually
     clean_message = message.replace("/nothink", "").strip()
 
-    # Build prompt
+    # Build prompt with last 4 history turns for context window economy
     full_prompt = f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n"
-
     for msg in history[-4:]:
         full_prompt += f"<|im_start|>{msg['role']}\n{msg['content']}<|im_end|>\n"
 
@@ -160,28 +261,29 @@ def generate_response(
         vs = get_vectorstore(backend)
         if vs:
             try:
-                # Build metadata filter: specific article takes priority over category
-                filter_dict = None
-                if article_filter and article_filter != "All articles in category":
-                    filter_dict = {"article_title": article_filter}
-                elif category_filter and category_filter != "All categories":
-                    filter_dict = {"article_category": category_filter}
-
-                docs = vs.similarity_search(clean_message, k=5, filter=filter_dict)
+                docs = _rag_search(
+                    vs, clean_message, k=5,
+                    article_filter=article_filter,
+                    category_filter=category_filter,
+                )
                 if docs:
                     context = "\n\n".join(
                         f"[Article: {doc.metadata.get('article_title', 'N/A')}] "
                         f"{doc.page_content[:700]}"
                         for doc in docs
                     )
-                    logger.info("RAG retrieved %d chunks (filter=%s)", len(docs), filter_dict)
+                    logger.info(
+                        "RAG: %d chunks injected (article=%r, cat=%r)",
+                        len(docs), article_filter, category_filter,
+                    )
                 else:
-                    logger.warning("RAG returned 0 chunks for filter=%s", filter_dict)
+                    logger.warning("RAG returned 0 chunks — answering without context")
             except Exception as exc:
                 logger.error("RAG retrieval failed: %s", exc)
 
-    # Append /nothink to suppress Qwen3 thinking if requested
-    nothink_suffix = " /nothink" if suppress_thinking else ""
+    # Qwen3 /nothink MUST appear on its own line at the very end of the user turn.
+    # A leading space (e.g. " /nothink") is NOT recognised by the tokeniser.
+    nothink_suffix = "\n/nothink" if suppress_thinking else ""
 
     if context:
         full_prompt += (
@@ -189,15 +291,17 @@ def generate_response(
             f"Question: {clean_message}{nothink_suffix}<|im_end|>\n"
         )
     else:
-        full_prompt += f"<|im_start|>user\n{clean_message}{nothink_suffix}<|im_end|>\n"
+        full_prompt += (
+            f"<|im_start|>user\n{clean_message}{nothink_suffix}<|im_end|>\n"
+        )
 
     full_prompt += "<|im_start|>assistant\n"
 
     # Sanitise generation params
-    max_tokens_val   = int(max_tokens)   if max_tokens   is not None else 900
-    temp_val         = float(temperature) if temperature  is not None else 0.65
-    top_p_val        = float(top_p)       if top_p        is not None else 0.9
-    rep_penalty_val  = float(repeat_penalty) if repeat_penalty is not None else 1.1
+    max_tokens_val  = int(max_tokens)        if max_tokens      is not None else 900
+    temp_val        = float(temperature)     if temperature     is not None else 0.65
+    top_p_val       = float(top_p)           if top_p           is not None else 0.9
+    rep_penalty_val = float(repeat_penalty)  if repeat_penalty  is not None else 1.1
 
     partial_text = ""
     for chunk in llm(
@@ -220,13 +324,15 @@ with gr.Blocks(title="Article Q&A model") as demo:
     gr.Markdown(
         "Qwen3.5-4B DoRA fine-tuned on 350+ articles on change from robertolofaro.com — "
         "experimental demo on CPU-only, to test embedding methods (takes a few minutes, "
-        "no selection for the category yet) — updated as of 2026-05-05"
+        "you can restrict by category, and then a specific article) — updated as of 2026-05-05"
     )
     gr.Markdown(
         "**NOTAM:** by querying this model you access the articles and metadata "
         "available on robertolofaro.com and GitHub.  "
         "Answers reflect the article corpus only — do not treat them as advice, "
-        "just expression of a position derived from material contained within the articles- if you want to read actual positions expressed within articles, you can read the articles (see the model repository for all the links the the available options)."
+        "just expression of a position derived from material contained within the articles. "
+        "If you want to read actual positions expressed within articles, you can read the articles "
+        "(see the model repository for all links to the available options)."
     )
     gr.Markdown(
         "If, after getting an answer, you want something tailored to your context, "
@@ -271,10 +377,10 @@ with gr.Blocks(title="Article Q&A model") as demo:
     )
 
     with gr.Accordion("Advanced Generation Parameters", open=False):
-        max_tokens    = gr.Slider(256, 2048, value=900,  step=64,   label="Max Tokens")
-        temperature   = gr.Slider(0.0,  1.0, value=0.65, step=0.05, label="Temperature")
-        top_p         = gr.Slider(0.0,  1.0, value=0.9,  step=0.05, label="Top-p")
-        repeat_penalty = gr.Slider(1.0, 2.0, value=1.1,  step=0.05, label="Repeat Penalty")
+        max_tokens     = gr.Slider(256, 2048, value=900,  step=64,   label="Max Tokens")
+        temperature    = gr.Slider(0.0,  1.0, value=0.65, step=0.05, label="Temperature")
+        top_p          = gr.Slider(0.0,  1.0, value=0.9,  step=0.05, label="Top-p")
+        repeat_penalty = gr.Slider(1.0,  2.0, value=1.1,  step=0.05, label="Repeat Penalty")
 
     gr.ChatInterface(
         fn=generate_response,