Upload 4 files

Dockerfile

@@ -18,14 +18,22 @@ ARG DDB_URL=https://huggingface.co/datasets/edgarodriguez/depth_alpha/resolve/ma
 ARG RAG_URL=https://huggingface.co/datasets/edgarodriguez/depth_rag/resolve/main/rag_chunks.parquet
 
 # Small local models served in-container by Ollama (no external inference).
-ARG CHAT_MODEL=qwen2.5:3b
+# 1.5b (not 3b) for ~2x faster CPU inference on the free Space; flip to qwen2.5:3b
+# via the DEPTH_LLM_MODEL env var (no rebuild) if you later add a GPU/CPU upgrade.
+ARG CHAT_MODEL=qwen2.5:1.5b
 ARG EMBED_MODEL=nomic-embed-text
 
+# OLLAMA_KEEP_ALIVE=-1 keeps the model resident (never unload) so the 2nd+ replies
+# skip the reload that caused the ~420s timeout; flash-attention + q8 KV cache trim
+# CPU work per token.
 ENV OLLAMA_MODELS=/app/.ollama \
     DEPTH_OLLAMA_URL=http://127.0.0.1:11434 \
     DEPTH_LLM_MODEL=${CHAT_MODEL} \
     DEPTH_EMBED_MODEL=${EMBED_MODEL} \
-    DEPTH_RAG_PARQUET=/app/rag_chunks.parquet
+    DEPTH_RAG_PARQUET=/app/rag_chunks.parquet \
+    OLLAMA_KEEP_ALIVE=-1 \
+    OLLAMA_FLASH_ATTENTION=1 \
+    OLLAMA_KV_CACHE_TYPE=q8_0
 
 # R packages. The base image pins a dated CRAN snapshot whose ellmer/shinychat
 # are too old -- that ellmer lacked the `params` argument, so the assistant failed

assistant.R

@@ -86,8 +86,9 @@ ollama_embed <- function(text, cfg = assistant_config()) {
 }
 
 # Embed the question (with the nomic query prefix) and return the top-k corpus
-# passages as a single context string with [source] tags for citation.
-rag_search <- function(con, query, cfg = assistant_config(), k = 5) {
+# passages as a single context string with [source] tags for citation. k kept
+# small (2) so the injected context stays short -> cheaper CPU prefill.
+rag_search <- function(con, query, cfg = assistant_config(), k = 2) {
   if (is.null(con) || !file.exists(cfg$rag_path)) return("")
   qvec <- ollama_embed(paste0("search_query: ", query), cfg)
   if (length(qvec) == 0) return("")
@@ -109,53 +110,46 @@ rag_search <- function(con, query, cfg = assistant_config(), k = 5) {
 }
 
 # --- Grounding text from WEB_SPEC --------------------------------------------
-# Walk the accordion (label + tooltip + source) into a compact data dictionary
-# the model can read. Stays in lockstep with the UI -- no separate codebook.
+# Walk the accordion into a COMPACT data dictionary the model can read: just the
+# layer labels (one per line), no verbose tooltip bodies. Keeps the system prompt
+# small so CPU prefill stays cheap. Stays in lockstep with the UI -- no codebook.
 assistant_indicator_catalog_txt <- function() {
   if (!exists("WEB_SPEC")) return("")
   out <- character(0)
   for (grp in WEB_SPEC$accordion) {
     layer_items <- Filter(function(it) identical(it$kind, "layer"), grp$items)
-    if (length(layer_items) == 0) next
-    out <- c(out, paste0("## ", grp$label))
-    for (it in layer_items) {
-      tip  <- it$tooltip
-      body <- if (is.list(tip)) paste(c(tip$body, tip$meta), collapse = " ")
-              else tip %||% ""
-      out <- c(out, sprintf("- %s: %s", it$label, body))
-    }
+    for (it in layer_items) out <- c(out, paste0("- ", it$label))
   }
-  paste(out, collapse = "\n")
+  paste(unique(out), collapse = "\n")
 }
 
 assistant_system_prompt <- function(catalog_txt) {
   paste0(
     "You are the research assistant for DEPTH, an interactive map of migration ",
-    "risk in Mexico built by the University of Nottingham Rights Lab. The map ",
-    "overlays humanitarian, violence and infrastructure indicators on a ",
-    "hierarchical grid of Mexico.\n\n",
-    "How to answer:\n",
-    "- For ANY number, ranking, count or aggregation you MUST call a tool. ",
-    "Never invent, estimate or recall figures from memory.\n",
-    "- When the user refers to 'the selection', 'this area', 'the area I drew', ",
-    "'here', or 'visible', call describe_selection first to learn the scope.\n",
-    "- Use the <context> document excerpts supplied with the user's message for ",
-    "background, definitions and qualitative explanation. Cite the document ",
-    "title in square brackets when you rely on it.\n",
-    "- If the data and documents do not cover the question, say so plainly. ",
-    "Do not speculate.\n",
-    "- Reply in the user's language (English or Spanish). Be concise: a few ",
-    "sentences or a short list.\n\n",
-    "Indicators available on the map (label: meaning and source):\n",
+    "risk in Mexico (University of Nottingham Rights Lab).\n",
+    "Rules:\n",
+    "- For any number, count, ranking or average, call a tool. Never invent figures.\n",
+    "- The user's current map scope and its totals are in <scope>; use them directly. ",
+    "Call a tool only for a different cut (a ranking, or another indicator).\n",
+    "- Use <context> document excerpts for definitions/background; cite the title in ",
+    "[brackets].\n",
+    "- If the data/documents don't cover it, say so. Reply in the user's language ",
+    "(English/Spanish), concisely.\n",
+    "Indicators on the map:\n",
     catalog_txt
   )
 }
 
-# Prepend retrieved corpus context to the raw question (the UI still shows the
-# raw question; only the model sees the augmented version).
-assistant_user_turn <- function(question, context_txt) {
-  if (is.null(context_txt) || !nzchar(context_txt)) return(question)
-  paste0("<context>\n", context_txt, "\n</context>\n\n", question)
+# Wrap the raw question with the current map scope and retrieved corpus context
+# (the UI still shows the raw question; only the model sees the augmented version).
+# Injecting <scope> here saves a whole tool round-trip on CPU.
+assistant_user_turn <- function(question, context_txt = "", scope_txt = "") {
+  parts <- character(0)
+  if (!is.null(scope_txt)   && nzchar(scope_txt))
+    parts <- c(parts, paste0("<scope>\n",   scope_txt,   "\n</scope>"))
+  if (!is.null(context_txt) && nzchar(context_txt))
+    parts <- c(parts, paste0("<context>\n", context_txt, "\n</context>"))
+  paste(c(parts, question), collapse = "\n\n")
 }
 
 # --- Tool back-ends (operate on a plain snapshot env + read-only con) ---------
@@ -308,20 +302,16 @@ make_assistant <- function(con, state, cfg = assistant_config()) {
 
   ind_enum <- ellmer::type_enum(names(.IND_LABELS), "Indicator column.")
 
-  list_indicators <- function() assistant_indicator_catalog_txt()
-  describe_selection <- function() assistant_selection_summary(state)
+  # Only the 3 data tools are registered. The indicator catalog lives in the
+  # system prompt and the current scope/totals are injected as <scope> each turn,
+  # so list_indicators / describe_selection are unnecessary -- fewer tool schemas
+  # mean less CPU prefill and easier tool choice for the 1.5B model.
   aggregate_indicator <- function(indicator, statistic = "mean")
     assistant_aggregate(con, state, indicator, statistic)
   top_municipalities <- function(indicator, n = 5)
     assistant_top_municipalities(con, state, indicator, n)
   count_points <- function(layer) assistant_count_points(state, layer)
 
-  chat$register_tool(ellmer::tool(list_indicators,
-    "List the indicators / data layers available in the DEPTH map, with their meaning and source. Call this when the user asks what data exists."))
-
-  chat$register_tool(ellmer::tool(describe_selection,
-    "Summarise the user's CURRENT map selection (drawn area, clicked municipality or state): how many municipalities and points of each type fall inside, plus key totals. Returns national totals if nothing is selected."))
-
   chat$register_tool(ellmer::tool(aggregate_indicator,
     "Aggregate one municipality indicator over the current selection (or nationally if nothing is selected). Use 'mean' for rate indicators and 'sum' for count indicators.",
     arguments = list(

depth_interactive_map_alpha.R

@@ -173,26 +173,26 @@ POBREZA_PAL  <- c("#fff7f3", "#feebe2", "#fbb4b9", "#f768a1",
 # Categorical palettes for the two non-numeric patp sub-types. Stable colour-blind
 # safe Set2-style hues; keys must match the truncated label list in setPATPType().
 PATP_RURAL_PROBLEM_PAL <- c(
-  "Sociales y económicos"                                       = "#66c2a5",
-  "Carencia o deficiencia en servicios básicos"                 = "#fc8d62",
-  "Caminos, transportes y telecomunicaciones"                        = "#8da0cb",
-  "Inseguridad, delincuencia y adicciones"                           = "#e78ac3",
-  "Carencia o deficiencia en la infraestructura educativa y de salud"= "#a6d854",
-  "Municipios sin localidades objetivo"                              = "#ffd92f",
-  "Sin problema"                                                     = "#e5c494",
+  "Sociales y económicos"                                       = "#4269D0",
+  "Carencia o deficiencia en servicios básicos"                 = "#EFB118",
+  "Inseguridad, delincuencia y adicciones"                           = "#FF725C",
+  "Caminos, transportes y telecomunicaciones"                        = "#FF8AB7",
+  "Carencia o deficiencia en la infraestructura educativa y de salud" = "#6CC5B0",
+  "Riesgos ambientales"                                              = "#3CA951",
   "Falta de información / Otro problema"                        = "#b3b3b3",
-  "Riesgos ambientales"                                              = "#1f78b4",
-  "Problema no definido"                                             = "#cab2d6"
+  "Problema no definido"                                             = "#cab2d6",
+  "Sin problema"                                                     = "#FBF9F1",
+  "Municipios sin localidades objetivo"                              = "#9C6B4E"
 )
 PATP_DELITO_PAL <- c(
-  "El patrimonio"                      = "#e41a1c",
-  "otros bienes"                       = "#377eb8",
-  "La vida y la Integridad corporal"   = "#4daf4a",
-  "múltiples bienes"               = "#984ea3",
-  "La familia"                         = "#ff7f00",
-  "Sin delito"                         = "#a6cee3",
-  "La libertad y la seguridad sexual"  = "#b15928",
-  "Libertad personal"                  = "#6a3d9a"
+  "El patrimonio"                      = "#4269D0",
+  "múltiples bienes"               = "#EFB118",
+  "otros bienes"                       = "#FF725C",
+  "La vida y la Integridad corporal"   = "#FF8AB7",
+  "La familia"                         = "#6CC5B0",
+  "La libertad y la seguridad sexual"  = "#3CA951",
+  "Libertad personal"                  = "#A463F2",
+  "Sin delito"                         = "#FBF9F1"
 )
 
 # Spanish-key -> English-display maps used by acc_legend_cats() so that the
@@ -231,10 +231,10 @@ PATP_DELITO_LABELS_EN <- c(
 # WEB_SQL_TEMPLATES$load_infra_pts maps the raw values to these buckets.
 INFRA_CATS <- c(
   "Shelter"              = "#3C6E71",   # forest green
-  "Support (NGO)"      = "#F5CB5C",   # ochre
-  "Support (GOV)"        = "#3A4B5C",   # slate
-  "Embassy / Consulate"  = "#4D194D",   # deep navy (international protection)
-  "Migration office"         = "#9E2A3A",   # brick red
+  "Support (NGO)"         = "#F5CB5C",   # ochre
+  "Support (GOV)"         = "#3A4B5C",   # slate
+  "Embassy / Consulate"   = "#4D194D",   # deep navy (international protection)
+  "Migration office"      = "#9E2A3A",   # brick red
   "Other"          = "#94a3b8"    # neutral slate (fallback)
 )
 
@@ -3921,9 +3921,12 @@ server <- function(input, output, session) {
           "Sorry -- the local model is unavailable right now.")
         return()
       }
-      ctx    <- tryCatch(rag_search(con, q, .ASSIST_CFG, k = 5),
+      ctx    <- tryCatch(rag_search(con, q, .ASSIST_CFG, k = 2),
                          error = function(e) "")
-      prompt <- assistant_user_turn(q, ctx)
+      # Inject the current map scope so the model skips a describe_selection round-trip.
+      scope  <- tryCatch(assistant_selection_summary(assist_state),
+                         error = function(e) "")
+      prompt <- assistant_user_turn(q, ctx, scope)
       tryCatch(
         shinychat::chat_append("assist_chat", chat$stream_async(prompt)),
         error = function(e) {

entrypoint.sh

@@ -42,5 +42,16 @@ ensure_model "${EMBED_MODEL}"
 echo "[entrypoint] ollama models available at runtime:"
 ollama list || true
 
+# Pre-warm both models so the FIRST user question hits a hot model -- the CPU load
+# of the weights into RAM happens now (at startup) instead of on that first reply.
+# Combined with OLLAMA_KEEP_ALIVE=-1 this removes the cold-load and reload latency.
+echo "[entrypoint] pre-warming ${CHAT_MODEL} (loads weights into RAM)..."
+curl -sf "http://${OLLAMA_HOST}/api/generate" \
+  -d "{\"model\":\"${CHAT_MODEL}\",\"prompt\":\"hi\",\"stream\":false,\"keep_alive\":-1}" \
+  >/dev/null 2>&1 && echo "[entrypoint] chat model warm." || echo "[entrypoint] chat warm skipped."
+curl -sf "http://${OLLAMA_HOST}/api/embed" \
+  -d "{\"model\":\"${EMBED_MODEL}\",\"input\":\"warm\"}" \
+  >/dev/null 2>&1 && echo "[entrypoint] embed model warm." || echo "[entrypoint] embed warm skipped."
+
 echo "[entrypoint] launching Shiny app..."
 exec Rscript depth_interactive_map_alpha.R