Spaces:

IOE-IFAD
/

test

Paused

test / _backend /reset_vdb.py

Automated build & deploy

e454383 verified about 1 year ago

2.49 kB

import os
from dotenv import load_dotenv
from utils.document_management import load_document, transcript_chunker, merge_chunks
from utils.vector_db import VectorDB
from langchain.schema import Document
from uuid import uuid4

# Load environment variables
load_dotenv()

# Initialize services
# Load all the transcriptions from the "transcription_data" directory
working_dir = os.path.join(os.getcwd(), "_backend","data" ,"honduras_data")
if os.getenv("ENVIRONMENT") == "production":
    working_dir = os.path.join(os.getcwd(), "data","honduras_data")
vector_db_name = "transcription"
vector_db_embedding_length = 1536

# First, collect all documents
all_documents = []
all_files = []

# Walk through the directory to find all *_transcribed.srt files
for root, dirs, files in os.walk(working_dir):
    for file in files:
        if file.endswith("transcribed.srt"):
            file_path = os.path.join(root, file)
            # clipping everything before honduras_data 
            relative_path = os.path.relpath(file_path, working_dir)
            all_files.append(relative_path)

for file_path in all_files:
    path = os.path.join(working_dir, file_path)
    
    docuemnt = load_document(path)

    # Chunk the document into smaller parts
    chunks = transcript_chunker(docuemnt, filename=file_path)

    merged_chunks = merge_chunks(chunks)

    # Create LangChain Document objects
    documents = [Document(page_content=chunk["text"], metadata={**chunk, "chunk_id": i, "id": str(uuid4()), "filename": file_path}) for i, chunk in enumerate(merged_chunks)]

    print(f"Processing {file_path} with {len(documents)} chunks")
    all_documents.extend(documents)


# Now create the vector database with all documents at once (this will delete and recreate)
if all_documents:
    if os.path.exists(os.path.join(working_dir, "chroma_db")):
        import shutil
        # Remove the existing vector database directory
        shutil.rmtree(os.path.join(working_dir, "chroma_db"))
        print("Removed existing vector database directory.")

        
    vector_db = VectorDB(
        collection_name=vector_db_name,
        embedding_length=vector_db_embedding_length,
        open_ai_api_key=os.getenv("OPENAI_API_KEY"),
        working_dir=working_dir,
        documents=all_documents
    )
    print(f"Created new vector database with {len(all_documents)} documents")
else:
    print("No documents found to process")

print("Vector database reset and data upload complete with ChromaDB.")

	import os
	from dotenv import load_dotenv
	from utils.document_management import load_document, transcript_chunker, merge_chunks
	from utils.vector_db import VectorDB
	from langchain.schema import Document
	from uuid import uuid4

	# Load environment variables
	load_dotenv()

	# Initialize services
	# Load all the transcriptions from the "transcription_data" directory
	working_dir = os.path.join(os.getcwd(), "_backend","data" ,"honduras_data")
	if os.getenv("ENVIRONMENT") == "production":
	working_dir = os.path.join(os.getcwd(), "data","honduras_data")
	vector_db_name = "transcription"
	vector_db_embedding_length = 1536

	# First, collect all documents
	all_documents = []
	all_files = []

	# Walk through the directory to find all *_transcribed.srt files
	for root, dirs, files in os.walk(working_dir):
	for file in files:
	if file.endswith("transcribed.srt"):
	file_path = os.path.join(root, file)
	# clipping everything before honduras_data
	relative_path = os.path.relpath(file_path, working_dir)
	all_files.append(relative_path)

	for file_path in all_files:
	path = os.path.join(working_dir, file_path)

	docuemnt = load_document(path)

	# Chunk the document into smaller parts
	chunks = transcript_chunker(docuemnt, filename=file_path)

	merged_chunks = merge_chunks(chunks)

	# Create LangChain Document objects
	documents = [Document(page_content=chunk["text"], metadata={**chunk, "chunk_id": i, "id": str(uuid4()), "filename": file_path}) for i, chunk in enumerate(merged_chunks)]

	print(f"Processing {file_path} with {len(documents)} chunks")
	all_documents.extend(documents)


	# Now create the vector database with all documents at once (this will delete and recreate)
	if all_documents:
	if os.path.exists(os.path.join(working_dir, "chroma_db")):
	import shutil
	# Remove the existing vector database directory
	shutil.rmtree(os.path.join(working_dir, "chroma_db"))
	print("Removed existing vector database directory.")


	vector_db = VectorDB(
	collection_name=vector_db_name,
	embedding_length=vector_db_embedding_length,
	open_ai_api_key=os.getenv("OPENAI_API_KEY"),
	working_dir=working_dir,
	documents=all_documents
	)
	print(f"Created new vector database with {len(all_documents)} documents")
	else:
	print("No documents found to process")

	print("Vector database reset and data upload complete with ChromaDB.")