jacob-mathison/doczyai-pipelines
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Activity

Merged in feature/faiss (pull request #680)

Browse Source 
Feature/faiss

* Migrate to FAISS from chroma

* Refactor safe_process_file to remove unnecessary garbage collection and update docstring for clarity

* Merged main into feature/faiss


Approved-by: Katon Minhas
This commit is contained in:
Alex Galarce
2025-08-26 16:54:55 +00:00
parent 94251faa03
commit 90bdc672d5
2 changed files with 43 additions and 388 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fieldExtraction/src/investment/main.py
+2 -13
View File
@@ -1,5 +1,4 @@
import concurrent.futures
import gc
import logging
import os
import random
@@ -10,10 +9,7 @@ import pandas as pd
random.seed(42)
import traceback
from constants.constants import Constants
from sentence_transformers import SentenceTransformer
import src.investment.file_processing as file_processing
import src.utils.io_utils as io_utils
@@ -21,13 +17,11 @@ from src import config
def safe_process_file(item, constants, run_timestamp):
"""Process a single file with error handling and resource cleanup.
"""Process a single file with comprehensive error handling.
This function wraps the main file processing logic to ensure that:
- Individual file failures don't crash the entire batch
- Resources are properly cleaned up after each file
- Detailed error information is captured for debugging
- Garbage collection is forced to help prevent file descriptor leaks
Args:
item: Tuple of (filename, contract_text) representing the file to process
@@ -41,16 +35,13 @@ def safe_process_file(item, constants, run_timestamp):
Note:
Always returns a DataFrame to maintain consistent output format for pd.concat.
Forces garbage collection after processing to help with resource cleanup,
particularly important for preventing file descriptor leaks in concurrent processing.
With FAISS migration, no special resource cleanup is needed.
"""
file_id = (
item[0] if isinstance(item, (list, tuple)) and len(item) > 0 else item
) # unpack file_id from item (passed in as a tuple below)
try:
result = file_processing.process_file(item, constants, run_timestamp)
# Force garbage collection after each file
gc.collect()
return result
except (
Exception
@@ -65,8 +56,6 @@ def safe_process_file(item, constants, run_timestamp):
logging.error(f"Full traceback:\n{full_traceback}")
# Force garbage collection on error too
gc.collect()
return pd.DataFrame(
[
{
fieldExtraction/src/investment/smart_chunking_funcs.py
+41 -375
View File
@@ -4,59 +4,35 @@ Smart Chunking Functions for Contract Processing
This module provides intelligent text chunking and semantic search capabilities
for contract processing. It includes:
- Vectorstore caching to prevent "too many open files" errors in concurrent processing
- Lightweight in-memory vectorstore operations using FAISS
- Multiple chunking strategies (hierarchical, OR, AND, regex)
- Semantic and keyword-based chunk retrieval
- Thread-safe operations for multi-threaded contract processing
- Monitoring and cleanup to prevent file descriptor leaks ("too many open files" errors)
- Optimized for ephemeral, per-contract vectorstore usage
Key Components:
- field_context(): Main function for semantic chunk retrieval with caching
- field_context(): Main function for semantic chunk retrieval using FAISS
- group_fields(): Groups fields based on retrieved text overlap
- Various chunk_*() functions: Different keyword matching strategies
Performance Notes:
- Uses module-level vectorstore caching to share embeddings across threads
- Reference counting ensures proper cleanup of temporary vectorstores
- Designed to handle concurrent processing without file descriptor leaks
- Emergency cleanup via atexit handler prevents resource leaks on process termination
File Descriptor Management:
- Monitors open file descriptors to prevent "too many open files" errors
- Implements reference counting for shared vectorstores
- Forces Chroma client connection cleanup via _client.reset()
- Includes retry logic for directory cleanup operations
- Logs detailed statistics every 10 files for monitoring
Troubleshooting:
- Check logs for "High number of open files" warnings
- Monitor cache hit/miss ratios for performance
- Watch for cleanup failure counts in cache statistics
- Emergency cleanup logs appear during process shutdown
- Uses FAISS for fast, in-memory similarity search
- No persistent storage - vectorstores are created and destroyed per contract
- Designed for concurrent processing without resource leaks
- Much lower memory footprint than persistent vectorstore solutions
"""
import atexit
import gc
import hashlib
import json
import logging
import os
import re
import shutil
import threading
import time
import weakref
from collections import defaultdict
from itertools import combinations
from typing import Any, Dict, List, Optional, Set, Tuple, Union
import chromadb
import nltk
import psutil
from constants.constants import Constants
from langchain_aws.embeddings.bedrock import BedrockEmbeddings
from langchain_chroma import Chroma
from langchain_community.retrievers import BM25Retriever
from langchain_community.vectorstores import FAISS
from langchain_text_splitters import RecursiveCharacterTextSplitter
from nltk.tokenize import word_tokenize
@@ -70,72 +46,6 @@ nltk.data.path.append(
)
)
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Enable Chroma reset functionality
try:
# Enable reset functionality to allow proper cleanup
chromadb.config.Settings(allow_reset=True)
logger.info("Chroma reset functionality enabled")
except Exception as e:
logger.warning(f"Could not enable Chroma reset: {e}")
# ===================================================================
# VECTORSTORE CACHING
# ===================================================================
# Thread-safe cache to prevent "too many open files" errors
_vectorstore_lock = threading.Lock()
_vectorstore_cache: Dict[str, Dict[str, Any]] = {}
# Track active vectorstores for cleanup
_active_vectorstores = weakref.WeakSet()
def _cleanup_all_vectorstores():
"""Emergency cleanup function for atexit"""
logger.info("Emergency cleanup starting...")
with _vectorstore_lock:
cleanup_count = 0
for cache_key in list(_vectorstore_cache.keys()):
try:
vectorstore_data = _vectorstore_cache[cache_key]
vectorstore = vectorstore_data["vectorstore"]
# Force close the vectorstore
if hasattr(vectorstore, "_client") and vectorstore._client:
vectorstore._client.reset()
if hasattr(vectorstore, "delete_collection"):
vectorstore.delete_collection()
# Clean up directory
chroma_filename = vectorstore_data.get("chroma_filename")
if os.path.exists(chroma_filename):
shutil.rmtree(chroma_filename, ignore_errors=True)
cleanup_count += 1
except Exception as e:
logger.warning(f"Error during emergency cleanup of {cache_key}: {e}")
_vectorstore_cache.clear()
logger.info(
f"Emergency cleanup completed. Cleaned up {cleanup_count} vectorstores."
)
# Register emergency cleanup
atexit.register(_cleanup_all_vectorstores)
# Monitoring counters
_cache_stats = {
"cache_hits": 0,
"cache_misses": 0,
"vectorstores_created": 0,
"vectorstores_cleaned": 0,
"cleanup_failures": 0,
"open_files_peak": 0,
}
# Setup bedrock
embeddings_model = BedrockEmbeddings(
@@ -156,72 +66,6 @@ TOP_K = 8
DATE_PATTERN = r"(?i)\bday of\b|\b(?:jan(?:uary)?|feb(?:ruary)?|mar(?:ch)?|apr(?:il)?|may|jun(?:e)?|jul(?:y)?|aug(?:ust)?|sep(?:tember)?|oct(?:ober)?|nov(?:ember)?|dec(?:ember)?)\b|\b(?:\d{1,2}[\/\.-]\d{1,2}[\/\.-](?:\d{2}|\d{4})|\d{4}[\/\.-]\d{1,2}[\/\.-]\d{1,2}|\d{8}|\d{2}\d{2}|\d{4}[\/\.-]\d{1,2}[\/\.-]\d{1,2}|\d{4}[\.]\d{1,2}[\.]\d{1,2})\b"
def get_open_file_count() -> int:
"""Get current number of open file descriptors for monitoring.
Uses Linux /proc filesystem when available, falls back to psutil.
Returns -1 if neither method works (shouldn't happen in normal operation).
Returns:
int: Number of open file descriptors, or -1 if unable to determine
"""
try: # Linux-specific: count files in /proc/PID/fd directory
return len(os.listdir(f"/proc/{os.getpid()}/fd"))
except:
# Fallback if /proc not available (e.g., on Windows or Mac)
try:
return len(psutil.Process().open_files())
except:
return -1
def log_cache_stats():
"""Log comprehensive cache statistics for monitoring and debugging.
Logs include:
- Cache hit/miss ratios (higher hit ratio = better performance)
- Vectorstore creation/cleanup counts (should be roughly equal)
- Current and peak open file descriptor counts
- Active cache size (number of cached vectorstores)
- Cleanup failure count (should remain at 0)
"""
open_files = get_open_file_count()
_cache_stats["open_files_peak"] = max(_cache_stats["open_files_peak"], open_files)
logger.info(
f"Cache Stats - Hits: {_cache_stats['cache_hits']}, "
f"Misses: {_cache_stats['cache_misses']}, "
f"Created: {_cache_stats['vectorstores_created']}, "
f"Cleaned: {_cache_stats['vectorstores_cleaned']}, "
f"Failures: {_cache_stats['cleanup_failures']}, "
f"Open Files: {open_files}, "
f"Peak Files: {_cache_stats['open_files_peak']}, "
f"Active Cache: {len(_vectorstore_cache)}"
)
def monitor_file_descriptors(filename: str, operation: str) -> int:
"""Monitor file descriptor usage and warn if approaching limits.
Args:
filename: Name of file being processed (for logging context)
operation: Operation being performed (e.g., "start", "vectorstore_create")
Returns:
int: Current number of open file descriptors
Note:
Logs warning if file descriptor count exceeds 800 (threshold for 1024 limit)
"""
open_files = get_open_file_count()
if open_files > 800: # Warning threshold
logger.warning(
f"High number of open files ({open_files}) detected during {operation} on {filename}"
)
return open_files
def parse_chunk(chunk: str) -> tuple[int, str]:
"""Extract chunk number and content from a chunk string
@@ -253,18 +97,6 @@ def find_overlap(prev_chunk, next_chunk, min_overlap=3):
return 0 # No overlap found
def get_vectorstore_cache_key(contract_text: str) -> str:
"""Generate a cache key for the vector store based on the contract text.
Args:
contract_text (str): The text of the contract to generate a key for.
Returns:
str: The generated cache key.
"""
return hashlib.md5(contract_text.encode()).hexdigest()[:16]
def stitch_chunks(chunks: list[str], overlap_size: int = CHUNK_OVERLAP) -> str:
"""Stitch together a list of chunks, with an overlap of `overlap_size` characters
@@ -310,48 +142,11 @@ def stitch_chunks(chunks: list[str], overlap_size: int = CHUNK_OVERLAP) -> str:
def field_context(
one_to_one_fields: FieldSet, constants: Constants, contract_text: str, filename: str
) -> Tuple[FieldSet, List[str]]:
"""Thread-safe field context extraction with vectorstore caching.
This function creates chunks from whole contract text and returns relevant chunks
for each field based on semantic search and keyword search.
This function is the core of the file descriptor leak prevention system.
It implements:
- Thread-safe vectorstore caching with reference counting
- Comprehensive cleanup with retries and forced garbage collection
- File descriptor monitoring throughout the process
- Emergency cleanup registration for process termination
Args:
one_to_one_fields: Fields to process for semantic chunking
constants: Configuration constants
contract_text: Full contract text to chunk and embed
filename: Filename for logging context
Returns:
Tuple of (updated fieldset with retrieved text, list of all chunks)
Raises:
Various exceptions from Chroma/embedding operations
Note:
This function is designed to be called concurrently from multiple threads.
The vectorstore cache prevents duplicate embedding work and the reference
counting ensures cleanup only happens when all threads are done.
"""
# === MONITORING SETUP ===
# Track processing time and file descriptor usage for this operation
start_time = time.time()
initial_files = monitor_file_descriptors(filename, "start_field_context")
"""Lightweight field context extraction using FAISS (in-memory only)."""
# Convert contract_text to lowercase for case-insensitive processing
contract_text = contract_text.lower()
# Generate unique cache key for this contract
cache_key = get_vectorstore_cache_key(contract_text)
chroma_filename = f"chroma_cache_{cache_key}"
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP
)
@@ -363,172 +158,43 @@ def field_context(
for idx, split in enumerate(splits)
]
vectorstore = None
try:
# === THREAD-SAFE CACHE ACCESS ===
# Multiple threads may try to create / access the same vectorstore
# Lock ensures only one thread creates it, others wait and reuse
with _vectorstore_lock:
if cache_key not in _vectorstore_cache:
# === VECTORSTORE CREATION ===
# This is the expensive operation we would like to cache
# Creates Chroma database with persistent storage for performance
logger.info(f"Creating new vectorstore for {filename} (cache miss)")
_cache_stats["cache_misses"] += 1
_cache_stats["vectorstores_created"] += 1
# Create lightweight FAISS vectorstore (in-memory only)
vectorstore = FAISS.from_texts(texts=splits, embedding=embeddings_model)
# Create vectorstore with persistence (for performance)
vectorstore = Chroma.from_texts(
texts=splits,
embedding=embeddings_model,
persist_directory=chroma_filename,
)
vectorstore_retriever = vectorstore.as_retriever(search_kwargs={"k": TOP_K})
keyword_retriever = BM25Retriever.from_texts(
splits, k=TOP_K, preprocess_func=word_tokenize
)
_vectorstore_cache[cache_key] = {
"vectorstore": vectorstore,
"splits": splits,
"chroma_filename": chroma_filename,
"ref_count": 0,
"created_at": time.time(),
}
for field in one_to_one_fields.fields:
if field.field_type == "smart_chunked":
questions = {}
questions[field.field_name] = field.get_prompt(constants)
question = json.dumps(questions)
# === EMERGENCY CLEANUP REGISTRATION ===
# Add to weak set so emergency cleanup can find it if process crashes
_active_vectorstores.add(vectorstore) # Track for emergency cleanup
# Convert field keywords to lowercase for case-insensitive processing
field_keywords = [keyword.lower() for keyword in field.keywords]
after_create_files = monitor_file_descriptors(
filename, "vectorstore_create"
)
logger.info(
f"Vectorstore created. Files {initial_files} -> {after_create_files}"
)
else:
# === CACHE HIT ===
# Reusing existing vectorstore saves significant time and resources
logger.info(f"Using cached vectorstore for {filename} (cache hit)")
_cache_stats["cache_hits"] += 1
# get relevant chunks based on semantic search and keyword search and merge them
semantic_relevant_docs = vectorstore_retriever.invoke(question)
keyword_relevant_docs = keyword_retriever.invoke(" ;".join(field.keywords))
semantic_relevant_chunks = [
doc.page_content for doc in semantic_relevant_docs
]
keyword_relevant_chunks = [
doc.page_content for doc in keyword_relevant_docs
]
if field.field_name == "CONTRACT_TITLE":
keyword_relevant_chunks = keyword_relevant_chunks + splits[:5]
field.retrieved_text = sorted(
list(set(semantic_relevant_chunks + keyword_relevant_chunks))
)
# === REFERENCE COUNTING ===
# Track how many threads are using this vectorstore
# Only cleanup when count reaches 0
_vectorstore_cache[cache_key]["ref_count"] += 1
vectorstore_data = _vectorstore_cache[cache_key]
# Use the cached/created vectorstore
vectorstore = vectorstore_data["vectorstore"]
splits = vectorstore_data["splits"]
vectorstore_retriever = vectorstore.as_retriever(search_kwargs={"k": TOP_K})
keyword_retriever = BM25Retriever.from_texts(
splits, k=TOP_K, preprocess_func=word_tokenize
)
for field in one_to_one_fields.fields:
if field.field_type == "smart_chunked":
questions = {}
questions[field.field_name] = field.get_prompt(constants)
question = json.dumps(questions)
# Convert field keywords to lowercase for case-insensitive processing
field_keywords = [keyword.lower() for keyword in field.keywords]
# get relevant chunks based on semantic search and keyword search and merge them
semantic_relevant_docs = vectorstore_retriever.invoke(question)
keyword_relevant_docs = keyword_retriever.invoke(
" ;".join(field.keywords)
)
semantic_relevant_chunks = [
doc.page_content for doc in semantic_relevant_docs
]
keyword_relevant_chunks = [
doc.page_content for doc in keyword_relevant_docs
]
if field.field_name == "CONTRACT_TITLE":
keyword_relevant_chunks = keyword_relevant_chunks + splits[:5]
field.retrieved_text = sorted(
list(set(semantic_relevant_chunks + keyword_relevant_chunks))
)
# Check if no context or full_context, if so, switch attribute
if len(field.retrieved_text) == 0 or len(field.retrieved_text) == len(
splits
):
field.field_type = "full_context"
finally:
# === CLEANUP WITH REFERENCE COUNTING ===
# This section is critical for preventing file descriptor leaks
with _vectorstore_lock:
if cache_key in _vectorstore_cache:
_vectorstore_cache[cache_key]["ref_count"] -= 1
# === ACTUAL CLEANUP ===
# Only cleanup when no other threads are using this vectorstore
if _vectorstore_cache[cache_key]["ref_count"] == 0:
logger.info(f"Cleaning up vectorstore for {filename}")
try:
vectorstore_to_cleanup = _vectorstore_cache[cache_key][
"vectorstore"
]
chroma_filename_to_cleanup = _vectorstore_cache[cache_key][
"chroma_filename"
]
# === FORCE CONNECTION CLEANUP ===
# Directory deletion can fail due to timing issues
# Retry with small delays to handle race conditions
if (
hasattr(vectorstore_to_cleanup, "_client")
and vectorstore_to_cleanup._client
):
vectorstore_to_cleanup._client.reset()
# Delete collection first
if hasattr(vectorstore_to_cleanup, "delete_collection"):
vectorstore_to_cleanup.delete_collection()
# === RETRY LOGIC ===
# Directory deletion can fail due to timing issues
# Retry with small delays to handle race conditions
# Clean up directory with retries
for attempt in range(3):
try:
if os.path.exists(chroma_filename_to_cleanup):
shutil.rmtree(
chroma_filename_to_cleanup, ignore_errors=True
)
break # Success
except OSError as e:
if attempt == 2: # Last attempt
logger.warning(
f"Failed to cleanup {chroma_filename_to_cleanup} after 3 attempts: {e}"
)
_cache_stats["cleanup_failures"] += 1
else:
time.sleep(0.1) # Brief delay before retry
# Force garbage collection
gc.collect()
_cache_stats["vectorstores_cleaned"] += 1
except Exception as e:
logger.warning(f"Error cleaning up {cache_key}: {e}")
_cache_stats["cleanup_failures"] += 1
finally:
# Always remove from cache
del _vectorstore_cache[cache_key]
end_time = time.time()
final_files = monitor_file_descriptors(filename, "completion")
processing_time = end_time - start_time
logger.info(
f"Completed {filename} in {processing_time:.2f}s. "
f"Files: {initial_files} -> {final_files}"
)
# Log stats every 10 files
if (_cache_stats["cache_hits"] + _cache_stats["cache_misses"]) % 10 == 0:
log_cache_stats()
# Fallback: if no context or all context, treat as full_context
if len(field.retrieved_text) == 0 or len(field.retrieved_text) == len(
splits
):
field.field_type = "full_context"
return one_to_one_fields, splits