| """ |
| ui/agents.py |
| ------------ |
| Heavy model / graph loading, done once per server process and shared by every |
| browser session. |
| |
| Streamlit gave this to us for free via @st.cache_resource. Gradio has no |
| equivalent, and it serves requests from a thread pool, so a naive "load if |
| None" would let two simultaneous first-visitors each start a ~2 GB model load. |
| Hence the explicit double-checked locking below: the lock is held across the |
| load, and a second caller blocks and then sees the finished object. |
| """ |
|
|
| import threading |
|
|
| _agents_lock = threading.Lock() |
| _agents = None |
|
|
| _chatbot_lock = threading.Lock() |
| _chatbot_models = None |
|
|
|
|
| def load_agents(): |
| """Import the compiled LangGraph agents. Importing the search graph also |
| loads the SPECTER reranker model at module import time (by design).""" |
| global _agents |
| if _agents is None: |
| with _agents_lock: |
| if _agents is None: |
| from app.modules.intent.graph import graph as intent_graph |
| from app.modules.search.graph import graph as search_graph |
| _agents = (intent_graph, search_graph) |
| return _agents |
|
|
|
|
| def warm_chatbot_models(): |
| """Pre-load the chatbot's embedding + cross-encoder models so the first |
| 'Chat it out' click doesn't pay the model-load cost. We instantiate the |
| exact models the chatbot uses (BAAI/bge-base-en-v1.5 + BAAI/bge-reranker-base), |
| warming the weights into the HF/torch cache. |
| |
| Both are pinned to CPU: this runs at boot, outside any ZeroGPU window, and |
| its whole job is to pull weights down — the GPU copy is made later, inside |
| ui.gpu.vectorize_on_gpu, from the same warmed cache.""" |
| global _chatbot_models |
| if _chatbot_models is None: |
| with _chatbot_lock: |
| if _chatbot_models is None: |
| from vectorizeer import get_embeddings |
| from langchain_community.cross_encoders import HuggingFaceCrossEncoder |
| embeddings = get_embeddings(device="cpu") |
| reranker = HuggingFaceCrossEncoder( |
| model_name="BAAI/bge-reranker-base", |
| model_kwargs={"device": "cpu"}, |
| ) |
| _chatbot_models = (embeddings, reranker) |
| return _chatbot_models |
|
|
|
|
| def warm_chatbot_models_async(): |
| """Kick the chatbot warm-up onto a daemon thread. |
| |
| The Streamlit app warmed BOTH model sets behind one blocking splash, which |
| meant nobody saw a usable page until ~2 GB of weights had downloaded. Only |
| the agents are needed to act on the very first click, so we block on those |
| and let the chatbot models finish in the background — they have until the |
| user has framed an intent, run a search, and picked a paper, which is far |
| longer than the load takes. ensure_chat_ready() calls warm_chatbot_models() |
| anyway, so if the thread hasn't finished it simply blocks on the same lock. |
| """ |
| threading.Thread(target=warm_chatbot_models, daemon=True).start() |
|
|
