Spaces:
Sleeping
Sleeping
File size: 6,196 Bytes
6a1ad16 e55cfd3 6a1ad16 e55cfd3 6a1ad16 e55cfd3 6a1ad16 e55cfd3 6a1ad16 e55cfd3 6a1ad16 e55cfd3 6a1ad16 e55cfd3 6a1ad16 e55cfd3 6a1ad16 e55cfd3 6a1ad16 e55cfd3 6a1ad16 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 |
import os
import chromadb
from sentence_transformers import SentenceTransformer
from loguru import logger
class SentenceTransformerEmbeddings:
def __init__(self, model_name: str = 'all-MiniLM-L6-v2'):
self.model = SentenceTransformer(model_name)
def __call__(self, input: list[str]) -> list[list[float]]:
embeddings = self.model.encode(input)
return embeddings.tolist()
def initialize_chromadb():
"""Initialize ChromaDB and load documents if needed"""
try:
# Set up paths
base_path = os.path.dirname(os.path.abspath(__file__))
doc_path = os.path.join(base_path, 'a2023-45.txt')
index_path = os.path.join(base_path, 'index.txt')
chroma_path = os.path.join(base_path, 'chroma_db')
# Check if required files exist
if not os.path.exists(doc_path):
logger.error(f"Document file not found at {doc_path}")
return False
if not os.path.exists(index_path):
logger.error(f"Index file not found at {index_path}")
return False
# Ensure ChromaDB directory exists
os.makedirs(chroma_path, exist_ok=True)
# Initialize ChromaDB
chroma_client = chromadb.PersistentClient(path=chroma_path)
embedding_function = SentenceTransformerEmbeddings()
# Check if collection exists and has content
collections = chroma_client.list_collections()
collection_exists = any(col.name == "legal_documents" for col in collections)
if collection_exists:
collection = chroma_client.get_collection(
name="legal_documents",
embedding_function=embedding_function
)
if collection.count() > 0:
logger.info("ChromaDB collection already exists and has content")
return True
# If we get here, we need to create or repopulate the collection
logger.info("Loading documents into ChromaDB...")
# Delete existing collection if it exists
if collection_exists:
chroma_client.delete_collection("legal_documents")
# Create new collection
collection = chroma_client.create_collection(
name="legal_documents",
embedding_function=embedding_function
)
# Read and process documents
with open(doc_path, 'r', encoding='utf-8') as f:
document = f.read().strip()
with open(index_path, 'r', encoding='utf-8') as f:
index_content = [line.strip() for line in f.readlines() if line.strip()]
# Process document into sections
sections = []
current_section = ""
current_title = ""
for line in document.split('\n'):
line = line.strip()
if any(index_line in line for index_line in index_content):
if current_section and current_title:
sections.append({
"title": current_title,
"content": current_section.strip()
})
current_title = line
current_section = ""
else:
if line:
current_section += line + "\n"
if current_section and current_title:
sections.append({
"title": current_title,
"content": current_section.strip()
})
# Prepare and add data to ChromaDB
if sections:
documents = []
metadatas = []
ids = []
for i, section in enumerate(sections):
if section["content"].strip():
documents.append(section["content"])
metadatas.append({
"title": section["title"],
"source": "a2023-45.txt",
"section_number": i + 1
})
ids.append(f"section_{i+1}")
collection.add(
documents=documents,
metadatas=metadatas,
ids=ids
)
logger.info(f"Successfully loaded {len(documents)} sections into ChromaDB")
return True
else:
logger.error("No valid sections found in document")
return False
except Exception as e:
logger.error(f"Error initializing ChromaDB: {str(e)}")
return False
def test_chromadb_content():
"""Test if ChromaDB has the required content"""
try:
# First ensure ChromaDB is initialized
if not initialize_chromadb():
return False
# Set up ChromaDB path
base_path = os.path.dirname(os.path.abspath(__file__))
chroma_path = os.path.join(base_path, 'chroma_db')
# Initialize ChromaDB
chroma_client = chromadb.PersistentClient(path=chroma_path)
# Get collection
collection = chroma_client.get_collection(
name="legal_documents",
embedding_function=SentenceTransformerEmbeddings()
)
# Check collection size
count = collection.count()
if count == 0:
logger.error("Collection is empty")
return False
logger.info(f"Found {count} documents in ChromaDB")
# Test query to verify content
test_results = collection.query(
query_texts=["What are the general provisions?"],
n_results=1
)
if not test_results['documents']:
logger.error("Test query returned no results")
return False
logger.info("ChromaDB content verification successful")
return True
except Exception as e:
logger.error(f"Error testing ChromaDB: {str(e)}")
return False
if __name__ == "__main__":
success = test_chromadb_content()
if success:
print("ChromaDB content verification successful")
else:
print("ChromaDB content verification failed") |