File size: 2,493 Bytes
e454383 | 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 | 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.") |