import logging
import numpy as np
import pandas as pd
from pathlib import Path
from typing import Optional, List, Dict  # **Added Import**
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, pipeline
import tabula
import pytesseract
from pdf2image import convert_from_path
import re
import faiss
import pickle
from tqdm import tqdm
import torch
from llama_index.core import Document, StorageContext, VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.storage.docstore import SimpleDocumentStore
from llama_index.core.storage.index_store import SimpleIndexStore
from llama_index.core.vector_stores import SimpleVectorStore
from langchain_huggingface import HuggingFaceEmbeddings

# Configure logging
logging.basicConfig(
    level=logging.DEBUG,  # Set to DEBUG for more detailed logs
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler()
    ]
)

# Initialize Sentence-BERT model and tokenizer (local model)
model_name = 'all-MiniLM-L6-v2'
logging.info(f"Loading SentenceTransformer model: {model_name}")
sentence_model = SentenceTransformer(model_name)
tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")

# Initialize HuggingFaceEmbeddings with the SentenceTransformer model
huggingface_embedding = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2')

# Define the directory containing your PDFs and storage
PDFS_DIRECTORY = Path('D:/Contracts')
STORAGE_PERSIST_DIR = Path("D:/Storage_test4")
FAISS_INDEX_PATH = STORAGE_PERSIST_DIR / "faiss_index.bin"
EMBEDDINGS_FILE = STORAGE_PERSIST_DIR / "document_embeddings.pkl"
DOC_METADATA_FILE = STORAGE_PERSIST_DIR / "document_metadata.pkl"

# Ensure the storage directory exists
STORAGE_PERSIST_DIR.mkdir(parents=True, exist_ok=True)

# Initialize variables to store embeddings and document metadata
all_embeddings = []
all_metadata = []

# Initialize summarizer (BART model)
logging.info("Initializing BART summarizer")
summarizer = pipeline("summarization", model="facebook/bart-large-cnn", device=0 if torch.cuda.is_available() else -1)

# Define the embedding dimension
dimension = sentence_model.get_sentence_embedding_dimension()

# Summarize text using BART summarizer with detailed debugging
def summarize_text(text, max_length=200, min_length=30):
    if len(text) < min_length:
        logging.warning("Text too short for summarization.")
        return ""

    try:
        if not text.strip():
            logging.warning("Text is empty or whitespace-only, skipping summarization.")
            return ""

        logging.debug(f"Text length: {len(text)}")
        logging.debug(f"Text content: {text[:100]}...")  # Log the first 100 characters of the text

        summary = summarizer(text, max_length=max_length, min_length=min_length, do_sample=False)[0]['summary_text']
        logging.info("Generated summary for text.")
        return summary
    except Exception as e:
        logging.error(f"Error generating summary: {e}")
        return ""

# Extract text from PDF and images
def extract_text_from_image(image):
    try:
        text = pytesseract.image_to_string(image)
        logging.info("Extracted text from image.")
        return text
    except Exception as e:
        logging.error(f"Error extracting text from image: {e}")
        return ""

def extract_text_from_pdf(pdf_path):
    try:
        logging.info(f"Converting PDF to images: {pdf_path}")
        images = convert_from_path(pdf_path)
        full_text = ""
        page_texts = []

        for i, image in enumerate(images):
            logging.info(f"Extracting text from page {i + 1}")
            text = extract_text_from_image(image)
            page_texts.append(text)

        for i, text in enumerate(page_texts):
            full_text += f"\n[Page {i + 1}]\n" + text

        logging.info(f"Extracted text from PDF: {pdf_path}")
        return full_text, len(page_texts)
    except Exception as e:
        logging.error(f"Error processing PDF {pdf_path}: {e}")
        return "", 0

def extract_tables_from_pdf(pdf_path):
    try:
        logging.info(f"Extracting tables from PDF: {pdf_path}")
        tables = tabula.read_pdf(pdf_path, pages='all', multiple_tables=True)
        if tables:
            logging.info(f"Extracted {len(tables)} tables.")
        else:
            logging.warning("No tables found.")
        return tables
    except Exception as e:
        logging.error(f"Error extracting tables from PDF {pdf_path}: {e}")
        return []

def extract_text_and_tables(pdf_path):
    text, num_pages = extract_text_from_pdf(pdf_path)
    tables = extract_tables_from_pdf(pdf_path)
    return text, tables, num_pages

# Function to embed text using SentenceTransformer
def embed_text(text: str) -> Optional[np.ndarray]:
    if not text.strip():
        logging.warning("Received empty text for embedding.")
        return None
    try:
        embedding = sentence_model.encode(text, convert_to_numpy=True).astype('float32')
        if embedding.shape[0] != dimension:
            logging.error(f"Embedding dimension mismatch: expected {dimension}, got {embedding.shape[0]}")
            return None
        logging.debug(f"Embedding shape: {embedding.shape}")
        return embedding
    except Exception as e:
        logging.error(f"Error embedding text: {e}")
        return None

# Function to embed table
def embed_table(table_df) -> Optional[np.ndarray]:
    try:
        table_str = table_df.to_string(index=False)
        table_embedding = embed_text(table_str)
        return table_embedding
    except Exception as e:
        logging.error(f"Error embedding table: {e}")
        return None

# Chunk text for embedding
def chunk_text(text, chunk_size=100, overlap=30):
    tokens = tokenizer.encode(text, truncation=True)
    total_tokens = len(tokens)
    chunks = []
    page_ranges = []

    start_idx = 0
    end_idx = chunk_size

    while start_idx < total_tokens:
        if end_idx > total_tokens:
            end_idx = total_tokens

        chunk_tokens = tokens[start_idx:end_idx]
        chunk_text_decoded = tokenizer.decode(chunk_tokens, skip_special_tokens=True)
        chunk_text_clean = chunk_text_decoded.strip()

        # Determine the page range for the current chunk
        page_matches = re.findall(r'\[Page (\d+)\]', chunk_text_clean)
        if page_matches:
            unique_pages = set(map(int, page_matches))
            chunk_start_page = min(unique_pages)
            chunk_end_page = max(unique_pages)
        else:
            chunk_start_page = (start_idx // chunk_size) + 1
            chunk_end_page = chunk_start_page

        page_ranges.append((chunk_start_page, chunk_end_page))
        chunks.append(chunk_text_clean)

        # Move start index forward with overlap
        start_idx += (chunk_size - overlap)
        end_idx = start_idx + chunk_size

    logging.info(f"Created {len(chunks)} text chunks with overlapping.")
    return chunks, page_ranges

# Function to save embeddings using FAISS for passage retrieval
def save_embeddings(text, text_embeddings, tables, table_embeddings, doc_id_prefix, faiss_index_obj, docstore, vector_store, index_store):
    documents = []

    # Ensure text is non-empty before summarizing
    if text and len(text.strip()) > 0:
        # Summarize the full text for Document-Level and Full Document Summary
        try:
            logging.info(f"Generating document-level summary for: {doc_id_prefix}")
            document_level_summary = summarize_text(text)
            logging.info(f"Document-level summary generated successfully for: {doc_id_prefix}")

            logging.info(f"Generating full document summary for: {doc_id_prefix}")
            full_document_summary = summarize_text(text, max_length=300)  # Longer summary for full document
            logging.info(f"Full document summary generated successfully for: {doc_id_prefix}")
        except Exception as e:
            logging.error(f"Error generating summary for {doc_id_prefix}: {e}")
            document_level_summary = "No valid summary could be generated."
            full_document_summary = "No valid full document summary could be generated."
    else:
        logging.warning(f"Empty or invalid text for {doc_id_prefix}. Skipping summary generation.")
        document_level_summary = "No valid summary available."
        full_document_summary = "No valid full document summary available."

    # Create and embed summaries if they exist
    document_summary_embedding = embed_text(document_level_summary)
    full_document_summary_embedding = embed_text(full_document_summary)

    if document_summary_embedding is not None:
        documents.append({
            "id": f"{doc_id_prefix}_document_summary",
            "text": document_level_summary,
            "embedding": document_summary_embedding,
            "extra_info": {"name": doc_id_prefix}
        })
        faiss_index_obj.add(document_summary_embedding.reshape(1, -1))
        vector_store.add([document_summary_embedding], [f"{doc_id_prefix}_document_summary"])
        index_store.add([f"{doc_id_prefix}_document_summary"])

    if full_document_summary_embedding is not None:
        documents.append({
            "id": f"{doc_id_prefix}_full_document_summary",
            "text": full_document_summary,
            "embedding": full_document_summary_embedding,
            "extra_info": {"name": doc_id_prefix}
        })
        faiss_index_obj.add(full_document_summary_embedding.reshape(1, -1))
        vector_store.add([full_document_summary_embedding], [f"{doc_id_prefix}_full_document_summary"])
        index_store.add([f"{doc_id_prefix}_full_document_summary"])

    # Process text chunks
    chunks, page_ranges = chunk_text(text)

    logging.info(f"Number of text chunks: {len(chunks)}")

    for i, embedding in enumerate(text_embeddings):
        if embedding is not None and i < len(chunks):
            doc_id = f"{doc_id_prefix}_text_chunk_{i}"
            documents.append({
                "id": doc_id,
                "text": chunks[i],
                "embedding": embedding,
                "extra_info": {"pages": f"{page_ranges[i][0]}-{page_ranges[i][1]}", "name": doc_id_prefix}
            })
            faiss_index_obj.add(embedding.reshape(1, -1))
            logging.debug(f"Adding embedding to vector store: {embedding}, {doc_id}")
            logging.debug(f"Embedding list: {[embedding]}")
            logging.debug(f"Doc ID list: {[doc_id]}")
            vector_store.add([embedding], [doc_id])
            index_store.add([doc_id])

    logging.info(f"Number of tables: {len(tables)}")
    for i, table_embedding in enumerate(table_embeddings):
        if table_embedding is not None and i < len(tables):
            doc_id = f"{doc_id_prefix}_table_{i}"
            table_text = tables[i].to_string(index=False)
            documents.append({
                "id": doc_id,
                "text": table_text,
                "embedding": table_embedding,
                "extra_info": {"pages": f"Table-{i + 1}", "name": doc_id_prefix}
            })
            faiss_index_obj.add(table_embedding.reshape(1, -1))
            logging.debug(f"Adding table embedding to vector store: {table_embedding}, {doc_id}")
            logging.debug(f"Table embedding list: {[table_embedding]}")
            logging.debug(f"Table doc ID list: {[doc_id]}")
            vector_store.add([table_embedding], [doc_id])
            index_store.add([doc_id])

    if not documents:
        logging.info(f"No documents to save for {doc_id_prefix}.")
        return

    # Append embeddings and metadata
    for doc in documents:
        all_embeddings.append(doc["embedding"])
        all_metadata.append({
            "id": doc["id"],
            "text": doc["text"],
            "extra_info": doc["extra_info"]
        })
        logging.info(f"Document {doc['id']} processed and added to FAISS index.")
        docstore.add_documents([Document(id_=doc["id"], text=doc["text"], extra_info=doc["extra_info"])])

def get_latest_versions(files):
    versioned_files = {}
    pattern = re.compile(r'_v(\d+\.\d+)')

    for file in files:
        match = pattern.search(file.stem)
        version = float(match.group(1)) if match else 1.0  # Default to v1.0 if no version
        base_name = pattern.sub('', file.stem)  # Remove version from base name

        # Track the latest version of the document
        if base_name not in versioned_files:
            versioned_files[base_name] = [(file, version)]
        else:
            versioned_files[base_name].append((file, version))

    latest_versions = []
    for base_name, file_versions in versioned_files.items():
        sorted_versions = sorted(file_versions, key=lambda x: x[1], reverse=True)
        latest_versions.append(sorted_versions[0][0])  # Keep only the latest version

    return latest_versions

def process_client_pdfs(client_dir, faiss_index_obj, docstore, vector_store, index_store):
    pdf_files = list(client_dir.glob("*.pdf"))
    latest_pdfs = get_latest_versions(pdf_files)  # Get the latest versions or valid variations

    for pdf in latest_pdfs:
        logging.info(f"Processing PDF: {pdf}")
        text, tables, num_pages = extract_text_and_tables(pdf)

        # Chunk, embed, and summarize as normal
        chunks, page_ranges = chunk_text(text)
        text_embeddings = [embed_text(chunk) for chunk in chunks]
        table_embeddings = [embed_table(table_df) for table_df in tables]

        # Save the embeddings and summaries with version info
        save_embeddings(text, text_embeddings, tables, table_embeddings, pdf.stem, faiss_index_obj, docstore, vector_store, index_store)

def main():
    client_dirs = [d for d in PDFS_DIRECTORY.iterdir() if d.is_dir()]
    logging.info(f"Found {len(client_dirs)} client directories to process.")

    # Initialize the storage context
    docstore = SimpleDocumentStore.from_persist_dir(persist_dir=STORAGE_PERSIST_DIR)
    vector_store = SimpleVectorStore.from_persist_dir(persist_dir=STORAGE_PERSIST_DIR)
    index_store = SimpleIndexStore.from_persist_dir(persist_dir=STORAGE_PERSIST_DIR)
    storage_context = StorageContext.from_defaults(
        docstore=docstore,
        vector_store=vector_store,
        index_store=index_store
    )

    # Initialize the FAISS index
    faiss_index = faiss.IndexFlatL2(dimension)

    for client_dir in tqdm(client_dirs, desc="Processing Clients"):
        client_name = client_dir.name
        logging.info(f"Processing client: {client_name}")

        # Process PDFs for this client, handling versions
        process_client_pdfs(client_dir, faiss_index, docstore, vector_store, index_store)

    # Save FAISS index and metadata
    try:
        logging.info(f"Saving FAISS index to {FAISS_INDEX_PATH}")
        faiss.write_index(faiss_index, str(FAISS_INDEX_PATH))
        logging.info(f"FAISS index saved successfully.")
    except Exception as e:
        logging.error(f"Error saving FAISS index: {e}")

    # Save embeddings and metadata
    try:
        with open(EMBEDDINGS_FILE, 'wb') as f:
            pickle.dump(all_embeddings, f)
            logging.info(f"Embeddings saved to {EMBEDDINGS_FILE}")
    except Exception as e:
        logging.error(f"Error saving embeddings: {e}")

    try:
        with open(DOC_METADATA_FILE, 'wb') as f:
            pickle.dump(all_metadata, f)
            logging.info(f"Document metadata saved to {DOC_METADATA_FILE}")
    except Exception as e:
        logging.error(f"Error saving document metadata: {e}")

    # Save the storage context
    storage_context.persist(persist_dir=STORAGE_PERSIST_DIR)

if __name__ == "__main__":
    main()


2024-09-27 12:11:43,344 - INFO - Loading SentenceTransformer model: all-MiniLM-L6-v2
2024-09-27 12:11:43,357 - INFO - Use pytorch device_name: cpu
2024-09-27 12:11:43,358 - INFO - Load pretrained SentenceTransformer: all-MiniLM-L6-v2
2024-09-27 12:11:46,553 - INFO - Use pytorch device_name: cpu
2024-09-27 12:11:46,555 - INFO - Load pretrained SentenceTransformer: sentence-transformers/all-MiniLM-L6-v2
2024-09-27 12:11:49,401 - INFO - Initializing BART summarizer
2024-09-27 12:11:55,011 - INFO - Found 2 client directories to process.
Processing Clients:   0%|          | 0/2 [00:00<?, ?it/s]2024-09-27 12:11:55,019 - INFO - Processing client: Absa
2024-09-27 12:11:55,021 - INFO - Processing PDF: D:\Contracts\Absa\ABSA 37 OCCUPATIONAL HEALTH AND SAFETY AGREEMENT MOBILISATION_v3.0.pdf
2024-09-27 12:11:55,022 - INFO - Converting PDF to images: D:\Contracts\Absa\ABSA 37 OCCUPATIONAL HEALTH AND SAFETY AGREEMENT MOBILISATION_v3.0.pdf
2024-09-27 12:12:00,495 - INFO - Extracting text from page 1
2024-09-27 12:12:02,011 - INFO - Extracted text from image.
2024-09-27 12:12:02,011 - INFO - Extracting text from page 2
2024-09-27 12:12:07,198 - INFO - Extracted text from image.
2024-09-27 12:12:07,198 - INFO - Extracting text from page 3
2024-09-27 12:12:12,823 - INFO - Extracted text from image.
2024-09-27 12:12:12,838 - INFO - Extracting text from page 4
2024-09-27 12:12:18,495 - INFO - Extracted text from image.
2024-09-27 12:12:18,495 - INFO - Extracting text from page 5
2024-09-27 12:12:23,919 - INFO - Extracted text from image.
2024-09-27 12:12:23,919 - INFO - Extracting text from page 6
2024-09-27 12:12:26,029 - INFO - Extracted text from image.
2024-09-27 12:12:26,029 - INFO - Extracted text from PDF: D:\Contracts\Absa\ABSA 37 OCCUPATIONAL HEALTH AND SAFETY AGREEMENT MOBILISATION_v3.0.pdf
2024-09-27 12:12:26,049 - INFO - Extracting tables from PDF: D:\Contracts\Absa\ABSA 37 OCCUPATIONAL HEALTH AND SAFETY AGREEMENT MOBILISATION_v3.0.pdf
2024-09-27 12:12:28,107 - WARNING - No tables found.
2024-09-27 12:12:28,161 - INFO - Created 8 text chunks with overlapping.
Error displaying widget
Error displaying widget
Error displaying widget
Error displaying widget
Error displaying widget
Error displaying widget
Error displaying widget
Error displaying widget
2024-09-27 12:12:30,632 - INFO - Generating document-level summary for: ABSA 37 OCCUPATIONAL HEALTH AND SAFETY AGREEMENT MOBILISATION_v3.0
2024-09-27 12:12:30,976 - ERROR - Error generating summary: index out of range in self
2024-09-27 12:12:31,017 - INFO - Document-level summary generated successfully for: ABSA 37 OCCUPATIONAL HEALTH AND SAFETY AGREEMENT MOBILISATION_v3.0
2024-09-27 12:12:31,018 - INFO - Generating full document summary for: ABSA 37 OCCUPATIONAL HEALTH AND SAFETY AGREEMENT MOBILISATION_v3.0
2024-09-27 12:12:31,083 - ERROR - Error generating summary: index out of range in self
2024-09-27 12:12:31,087 - INFO - Full document summary generated successfully for: ABSA 37 OCCUPATIONAL HEALTH AND SAFETY AGREEMENT MOBILISATION_v3.0
2024-09-27 12:12:31,089 - WARNING - Received empty text for embedding.
2024-09-27 12:12:31,090 - WARNING - Received empty text for embedding.
2024-09-27 12:12:31,141 - INFO - Created 8 text chunks with overlapping.
2024-09-27 12:12:31,143 - INFO - Number of text chunks: 8
Processing Clients:   0%|          | 0/2 [00:36<?, ?it/s]
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[55], line 388
    385     storage_context.persist(persist_dir=STORAGE_PERSIST_DIR)
    387 if __name__ == "__main__":
--> 388     main()

Cell In[55], line 359, in main()
    356     logging.info(f"Processing client: {client_name}")
    358     # Process PDFs for this client, handling versions
--> 359     process_client_pdfs(client_dir, faiss_index, docstore, vector_store, index_store)
    361 # Save FAISS index and metadata
    362 try:

Cell In[55], line 335, in process_client_pdfs(client_dir, faiss_index_obj, docstore, vector_store, index_store)
    332 table_embeddings = [embed_table(table_df) for table_df in tables]
    334 # Save the embeddings and summaries with version info
--> 335 save_embeddings(text, text_embeddings, tables, table_embeddings, pdf.stem, faiss_index_obj, docstore, vector_store, index_store)

Cell In[55], line 263, in save_embeddings(text, text_embeddings, tables, table_embeddings, doc_id_prefix, faiss_index_obj, docstore, vector_store, index_store)
    261         logging.debug(f"Embedding list: {[embedding]}")
    262         logging.debug(f"Doc ID list: {[doc_id]}")
--> 263         vector_store.add([embedding], [doc_id])
    264         index_store.add([doc_id])
    266 logging.info(f"Number of tables: {len(tables)}")

TypeError: SimpleVectorStore.add() takes 2 positional arguments but 3 were given


Its still not helping in any way focus brother