diff --git a/CHANGELOG.md b/CHANGELOG.md index 18c2b82c..7ba6f7ff 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,9 @@ - Handled multi-word property keywords (e.g., _thermal conductivity_) for accurate Scopus search, uniform filename handling (`thermal conductivity` resolves to `thermal_conductivity_metadata.csv` or similar) and restoring the original form `thermal conductivity` in the data extraction RAG search query instead of `thermal_conductivity`. This fix is associated with [#5](https://github.com/slimeslab/ComProScanner/pull/5) and contributed by [@WilmerGaspar](https://github.com/WilmerGaspar). +- Previously, a new `MultiModelEmbeddings` instance (and thus a fresh copy of the PhysBERT model) was loaded onto the GPU for every paper processed, because `RAGTool → VectorDatabaseManager → MultiModelEmbeddings` were all re-instantiated per paper. After certain number of papers this exhausted VRAM with `cudaErrorMemoryAllocation` (Refer to issue [#6](https://github.com/slimeslab/ComProScanner/issues/6)). This fix introduces a class-level `_hf_model_cache` dict on MultiModelEmbeddings so the tokenizer and model are loaded onto the GPU exactly once and shared as references across all subsequent instances. Also explicitly delete intermediate CUDA tensors and call `torch.cuda.empty_cache()` after each embedding call to prevent activation memory from accumulating within a paper's processing. Added the same cache flush in `VectorDatabaseManager.create_database` and `query_database` after `gc.collect()`. This fix is associated with PR [#7](https://github.com/slimeslab/ComProScanner/pull/7). + + --- # 2026.05.19 diff --git a/docs/about/changelog.md b/docs/about/changelog.md index f79a0060..3317088c 100644 --- a/docs/about/changelog.md +++ b/docs/about/changelog.md @@ -10,6 +10,8 @@ - Handled multi-word property keywords (e.g., _thermal conductivity_) for accurate Scopus search, uniform filename handling (`thermal conductivity` resolves to `thermal_conductivity_metadata.csv` or similar) and restoring the original form `thermal conductivity` in the data extraction RAG search query instead of `thermal_conductivity`. This fix is associated with [#5](https://github.com/slimeslab/ComProScanner/pull/5) and contributed by [@WilmerGaspar](https://github.com/WilmerGaspar). +- Previously, a new `MultiModelEmbeddings` instance (and thus a fresh copy of the PhysBERT model) was loaded onto the GPU for every paper processed, because `RAGTool → VectorDatabaseManager → MultiModelEmbeddings` were all re-instantiated per paper. After certain number of papers this exhausted VRAM with `cudaErrorMemoryAllocation` (Refer to issue [#6](https://github.com/slimeslab/ComProScanner/issues/6)). This fix introduces a class-level `_hf_model_cache` dict on MultiModelEmbeddings so the tokenizer and model are loaded onto the GPU exactly once and shared as references across all subsequent instances. Also explicitly delete intermediate CUDA tensors and call `torch.cuda.empty_cache()` after each embedding call to prevent activation memory from accumulating within a paper's processing. Added the same cache flush in `VectorDatabaseManager.create_database` and `query_database` after `gc.collect()`. This fix is associated with PR [#7](https://github.com/slimeslab/ComProScanner/pull/7). + --- # 2026.05.19 diff --git a/examples/vlm_piezo_test/vlm_test_example.py b/examples/vlm_piezo_test/vlm_test_example.py index 77a264c2..1583e446 100644 --- a/examples/vlm_piezo_test/vlm_test_example.py +++ b/examples/vlm_piezo_test/vlm_test_example.py @@ -204,7 +204,7 @@ test_doi_list_file="random_dois_for_vlm_test.txt", is_extract_synthesis_data=False, # For this test, we are only evaluating the composition-property extraction capability of the VLM, so we set this to False to save time and cost. model="deepseek/deepseek-chat", - vlm_model="google/gemini-3-flash-preview", + vlm_model="gemini/gemini-3-flash-preview", output_log_folder="vlm_piezo_test/model-logs/logs/google/gemini-3-flash-preview", task_output_folder="vlm_piezo_test/model-logs/task_outputs/google/gemini-3-flash-preview", materials_data_identifier_query="Is there any ceramic, composite, or crystal material with its chemical composition or doping data, and corresponding d33 piezoelectric coefficient value (in pC/N or pm/V units) mentioned in the text of this paper? Give one word answer - either 'yes' or 'no'. Only answer 'yes' if ALL of the following criteria are met: (1) The material is specifically a ceramic, composite, doped, or crystal, or different environments of materials (exclude all polymers including PVDF, PLLA, and similar), (2) A numerical d33 value with units pC/N or pm/V is explicitly stated, associated with either a specific material composition/environment or a doping concentration variable (e.g. x mol%, at%) — note that in figures/graphs, d33 values plotted against doping concentrations or composition variables also count as relevant data.", diff --git a/src/comproscanner/utils/database_manager.py b/src/comproscanner/utils/database_manager.py index bbcd266e..cdf5c08b 100644 --- a/src/comproscanner/utils/database_manager.py +++ b/src/comproscanner/utils/database_manager.py @@ -255,6 +255,12 @@ def create_database(self, db_name: str, article_text: str): self.client.clear_system_cache() del vectordb gc.collect() + try: + import torch + if torch.cuda.is_available(): + torch.cuda.empty_cache() + except ImportError: + pass logger.info(f"Vector database auto-persisted at {db_location}") @@ -274,6 +280,12 @@ def query_database(self, db_name: str, query: str, top_k: int = 5): self.client.clear_system_cache() del vectordb gc.collect() + try: + import torch + if torch.cuda.is_available(): + torch.cuda.empty_cache() + except ImportError: + pass logger.info(f"Retrieved {len(results)} results from {db_name}") return results diff --git a/src/comproscanner/utils/embeddings.py b/src/comproscanner/utils/embeddings.py index b50522f5..f499aa33 100644 --- a/src/comproscanner/utils/embeddings.py +++ b/src/comproscanner/utils/embeddings.py @@ -44,6 +44,10 @@ class MultiModelEmbeddings(Embeddings): """Embeddings class supporting multiple model types optimized for processing pre-chunked articles""" + # Class-level cache: model_name -> {"tokenizer": ..., "model": ..., "device": ...} + # Ensures the HuggingFace model is loaded onto GPU only once regardless of how many instances are created across papers. + _hf_model_cache: dict = {} + def __init__(self, rag_config: Any): """ Args: @@ -79,17 +83,32 @@ def _determine_model_type(self, model_name: str) -> str: return "openai" def _init_huggingface(self): - """Initialize HuggingFace model (works with any transformer model)""" - # Strip the prefix "huggingface:" from the model name + """Initialize HuggingFace model, reusing a cached instance if already loaded.""" model_name = self.rag_config.embedding_model if model_name.startswith("huggingface:"): model_name = model_name[len("huggingface:") :] - self.tokenizer = AutoTokenizer.from_pretrained(model_name) - self.model = AutoModel.from_pretrained(model_name) - self.device = "cuda" if torch.cuda.is_available() else "cpu" - self.model.to(self.device) - self.model.eval() + if model_name not in MultiModelEmbeddings._hf_model_cache: + logger.info( + f"Loading HuggingFace embedding model '{model_name}' onto device for the first time." + ) + tokenizer = AutoTokenizer.from_pretrained(model_name) + model = AutoModel.from_pretrained(model_name) + device = "cuda" if torch.cuda.is_available() else "cpu" + model.to(device) + model.eval() + MultiModelEmbeddings._hf_model_cache[model_name] = { + "tokenizer": tokenizer, + "model": model, + "device": device, + } + else: + logger.debug(f"Reusing cached HuggingFace embedding model '{model_name}'.") + + cached = MultiModelEmbeddings._hf_model_cache[model_name] + self.tokenizer = cached["tokenizer"] + self.model = cached["model"] + self.device = cached["device"] def _init_sentence_transformers(self): """Initialize SentenceTransformers model""" @@ -176,9 +195,20 @@ def _embed_document_huggingface(self, text: str) -> List[float]: sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1) sum_mask = torch.sum(input_mask_expanded, 1) sum_mask = torch.clamp(sum_mask, min=1e-9) - embedding = (sum_embeddings / sum_mask).cpu().numpy()[0] + embedding = (sum_embeddings / sum_mask).cpu().numpy()[0].tolist() + + del ( + inputs, + outputs, + token_embeddings, + input_mask_expanded, + sum_embeddings, + sum_mask, + ) + if self.device == "cuda": + torch.cuda.empty_cache() - return embedding.tolist() + return embedding def _embed_document_sentence_transformers(self, text: str) -> List[float]: """SentenceTransformers document embedding implementation for a single chunk"""