Spaces:
Sleeping
Sleeping
File size: 3,681 Bytes
6702977 |
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 |
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 load_documents():
"""Load and process documents into ChromaDB"""
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')
# Ensure ChromaDB directory exists
os.makedirs(chroma_path, exist_ok=True)
logger.info(f"Loading documents from {doc_path} and {index_path}")
# Initialize ChromaDB
chroma_client = chromadb.PersistentClient(path=chroma_path)
embedding_function = SentenceTransformerEmbeddings()
# Create new collection (delete if exists)
if "legal_documents" in [col.name for col in chroma_client.list_collections()]:
chroma_client.delete_collection("legal_documents")
collection = chroma_client.create_collection(
name="legal_documents",
embedding_function=embedding_function
)
# Read and validate files
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"
# Add final section
if current_section and current_title:
sections.append({
"title": current_title,
"content": current_section.strip()
})
# Prepare data for ChromaDB
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}")
# Add to ChromaDB
collection.add(
documents=documents,
metadatas=metadatas,
ids=ids
)
logger.info(f"Successfully loaded {len(documents)} sections into ChromaDB")
return True
except Exception as e:
logger.error(f"Error loading documents: {str(e)}")
return False
if __name__ == "__main__":
success = load_documents()
if success:
print("Documents successfully loaded into ChromaDB")
else:
print("Failed to load documents into ChromaDB") |