diff --git a/backend/app/rag/graph_retriever.py b/backend/app/rag/graph_retriever.py index 39841aa1..2ff1eca2 100644 --- a/backend/app/rag/graph_retriever.py +++ b/backend/app/rag/graph_retriever.py @@ -96,7 +96,15 @@ def get_entity_context( "pages": set(), }, ) - existing["weight"] = int(existing["weight"]) + int(edge.get("weight", 1)) + # Safely extract the raw weight, defaulting to 1 if missing or invalid + raw_weight = edge.get("weight", 1) + try: + # Convert to float first to handle decimals safely, then round/truncate if needed + weight_value = int(float(raw_weight)) if raw_weight is not None else 1 + except (ValueError, TypeError): + weight_value = 1 + + existing["weight"] += weight_value existing["pages"].update(edge.get("pages", [])) except Exception as exc: logger.warning("GraphRAG context retrieval failed: %s", exc) @@ -105,9 +113,10 @@ def get_entity_context( if not relationships: return "" + # Clean sorting without redundant typecasting overhead ranked = sorted( relationships.values(), - key=lambda item: int(item["weight"]), + key=lambda item: item["weight"], reverse=True, )[: settings.GRAPH_MAX_RELATIONSHIPS] diff --git a/backend/app/rag/reranker.py b/backend/app/rag/reranker.py index f43116ba..33c80412 100644 --- a/backend/app/rag/reranker.py +++ b/backend/app/rag/reranker.py @@ -31,13 +31,24 @@ def __init__(self, model_name: Optional[str] = None, device: Optional[str] = Non # Lazy-load the model when needed to avoid long startup times def _load_model(self) -> CrossEncoder: - """Lazy-load the cross-encoder model.""" + """Lazy-load the cross-encoder model with auto-precision if CUDA is available.""" if self._model is None: - logger.info(f"Loading reranker: {self.model_name}") + import torch + + # Detect device fallback if not explicitly set + current_device = self.device or ("cuda" if torch.cuda.is_available() else "cpu") + logger.info(f"Loading reranker: {self.model_name} on {current_device}") + + # Optimization: Use float16 on CUDA devices to slash VRAM footprint and speed up inference + model_kwargs = {} + if "cuda" in current_device: + model_kwargs["torch_dtype"] = torch.float16 + self._model = CrossEncoder( self.model_name, max_length=512, - device=self.device + device=current_device, + **model_kwargs ) logger.info("Reranker loaded successfully") return self._model @@ -49,6 +60,7 @@ def rerank( documents: List[Dict[str, Any]], top_k: int = 5, text_key: str = "text", + batch_size: int = 32, ) -> List[Dict[str, Any]]: """ Rerank documents based on relevance to the query. @@ -58,6 +70,7 @@ def rerank( documents: List of document dicts (must contain text_key field). top_k: Number of top documents to return after reranking. text_key: Key in document dict that holds the text content. + batch_size: Number of pairs to process simultaneously to prevent memory exhaustion. Returns: List of reranked documents (same dicts, but sorted by relevance). @@ -70,8 +83,8 @@ def rerank( # Prepare query-document pairs pairs = [(query, doc[text_key]) for doc in documents] - # Get relevance scores - scores = model.predict(pairs) + # Get relevance scores (utilizing batch_size to prevent OOM errors) + scores = model.predict(pairs, batch_size=batch_size) # Pair scores with documents and sort in descending order scored = list(zip(scores, documents))