correct_answer_generation/create_database.py

import fitz  # PyMuPDF
import re
import chromadb
import sys
import os
import uuid
sys.path.append(os.path.dirname(os.path.dirname(__file__)))
from all_models import models

def clean_text(text):
    # Keep only letters, numbers, punctuation, whitespace, and newlines
    cleaned_text = re.sub(r"[^a-zA-Z0-9\s.,!?;:'\"()\-]", "", text)
    return cleaned_text

def extract_text_from_pdf(pdf_path):
    text = ""
    with fitz.open(pdf_path) as doc:
        for page in doc:
            page_text = page.get_text()
            cleaned_text = clean_text(page_text)
            text += cleaned_text
    return text

def clean_data(text):
    cleaned_text = re.sub(r'\n{2,}', '. \n', text)  # Replace multiple newlines with a single newline
    cleaned_text = re.sub(r' {2,}', '. \n', cleaned_text)  # Replace multiple spaces with a newline
    return cleaned_text.strip()  # Strip leading/trailing whitespace

def combine_list(strings):
    combined_list = [] 
    current_combined = ""
    for string in strings:
        word_count = len(string.split())
        
        if len(current_combined.split()) < 20:
            current_combined += " " + string.strip()  # Adding space before new string
            
        # If the combined string reaches at least 20 words, add it to the final list
        if len(current_combined.split()) >= 20:
            combined_list.append(current_combined)  # Strip to remove leading/trailing whitespace
            current_combined = ""  # Reset for the next round
    if current_combined:
        combined_list.append(current_combined.strip())
    return combined_list

def create_databse(data, name):
    # Initialize the Persistent Client
    client = chromadb.PersistentClient(path="correct_answer_generation/chroma_db")
    
    collection_names = client.list_collections()
    if name in collection_names:
        client.delete_collection(name)  # Delete the old collection
    
    # Create a Collection
    collection = client.create_collection(name)

    # Generate embeddings using the singleton model
    embeddings = models.similarity_model.encode(data, batch_size=32, convert_to_tensor=True)

    # Create documents and add them to the collection
    unique_id = [str(uuid.uuid4()) for _ in range(len(embeddings))]
    
    collection.add(
        documents=data,
        ids=unique_id
    )

def create_database_main(path):
    pdf_path = path
    pdf_text = extract_text_from_pdf(pdf_path)
    data = clean_data(pdf_text)
    data = data.split('. \n')
    for i in range(len(data)):
        data[i] = re.sub(r' \n', ' ', data[i])
        data[i] = re.sub(r'\s+', ' ', data[i])
    data = [text for text in data if len(text) >= 2]
    data = combine_list(data)
    
    path = path.replace("/", "_")
    create_databse(data, path)