/usr/local/lib/python3.10/dist-packages/transformers/tokenization_utils_base.py:1601: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884
  warnings.warn(
ERROR:__main__:Error during document splitting: TextSplitter.__init__() got an unexpected keyword argument 'tokenizer'
ERROR:__main__:RAG pipeline failed: TextSplitter.__init__() got an unexpected keyword argument 'tokenizer'
import logging
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple, Union
import multiprocessing
import sys
import torch
import datasets
from datasets import Dataset
from langchain.docstore.document import Document as LangchainDocument
from langchain.text_splitter import RecursiveCharacterTextSplitter
from transformers import AutoTokenizer
from langchain.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores.utils import DistanceStrategy
import pandas as pd
from tqdm import tqdm
import numpy as np
import pacmap
import plotly.express as px
import plotly.graph_objects as go

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s %(levelname)s:%(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger(__name__)


@dataclass
class DataLoader:
    dataset_name: str
    split: str = "train"
    dataset: Optional[Dataset] = field(init=False, default=None)

    def load_dataset(self) -> None:
        """
        Load the dataset using Hugging Face datasets library.
        """
        try:
            logger.info(f"Loading dataset '{self.dataset_name}' with split '{self.split}'.")
            self.dataset = datasets.load_dataset(self.dataset_name, split=self.split)
            logger.info("Dataset loaded successfully.")
        except Exception as e:
            logger.error(f"Failed to load dataset '{self.dataset_name}': {e}")
            raise

    def identify_text_columns(self) -> List[str]:
        """
        Identify columns in the dataset that contain text data.
        """
        if self.dataset is None:
            logger.error("Dataset not loaded. Call load_dataset() first.")
            raise ValueError("Dataset not loaded.")

        text_columns = [column for column, dtype in self.dataset.features.items() if dtype == datasets.Value('string')]
        if not text_columns:
            logger.warning("No text columns found in the dataset.")
        else:
            logger.info(f"Identified text columns: {text_columns}")
        return text_columns


@dataclass
class DocumentProcessor:
    text_columns: List[str]
    chunk_size: int = 512
    chunk_overlap: int = 50
    tokenizer_name: str = "bert-base-uncased"
    multi_process: bool = True
    knowledge_base: List[LangchainDocument] = field(init=False, default_factory=list)
    docs_processed: List[LangchainDocument] = field(init=False, default_factory=list)

    def create_documents(self, dataset: Dataset) -> List[LangchainDocument]:
        """
        Convert dataset entries to LangChain documents.
        """
        logger.info("Converting dataset entries to LangChain documents.")
        documents = []
        for record in tqdm(dataset, desc="Creating documents"):
            # Concatenate all text columns
            page_content = ' '.join([str(record[col]) for col in self.text_columns if record[col]])
            metadata = {col: record[col] for col in self.text_columns if record[col]}
            documents.append(LangchainDocument(page_content=page_content, metadata=metadata))
        logger.info(f"Created {len(documents)} documents.")
        return documents

    def split_documents(self, documents: List[LangchainDocument]) -> List[LangchainDocument]:
        """
        Split documents into smaller chunks.
        """
        logger.info("Splitting documents into smaller chunks.")
        try:
            tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name)
            text_splitter = RecursiveCharacterTextSplitter(
                tokenizer=tokenizer,
                chunk_size=self.chunk_size,
                chunk_overlap=self.chunk_overlap,
                separators=["\n\n", "\n", " ", ""]
            )

            if self.multi_process:
                logger.info("Using multiprocessing for splitting documents.")
                with multiprocessing.Pool() as pool:
                    results = list(tqdm(pool.imap(text_splitter.split_documents, [documents[i:i + 1000] for i in range(0, len(documents), 1000)]), total=(len(documents) // 1000) + 1, desc="Splitting documents"))
                split_docs = [doc for sublist in results for doc in sublist]
            else:
                split_docs = []
                for doc in tqdm(documents, desc="Splitting documents"):
                    split_docs.extend(text_splitter.split_documents([doc]))
            logger.info(f"Split into {len(split_docs)} document chunks.")
            return split_docs
        except Exception as e:
            logger.error(f"Error during document splitting: {e}")
            raise

    def remove_duplicates(self, documents: List[LangchainDocument]) -> List[LangchainDocument]:
        """
        Remove duplicate documents based on page content.
        """
        logger.info("Removing duplicate documents.")
        unique_texts = set()
        unique_docs = []
        for doc in documents:
            if doc.page_content not in unique_texts:
                unique_texts.add(doc.page_content)
                unique_docs.append(doc)
        logger.info(f"Reduced to {len(unique_docs)} unique documents.")
        return unique_docs

    def process_documents(self, dataset: Dataset) -> None:
        """
        Full pipeline to process documents.
        """
        documents = self.create_documents(dataset)
        split_docs = self.split_documents(documents)
        unique_docs = self.remove_duplicates(split_docs)
        self.docs_processed = unique_docs
        logger.info(f"Document processing complete. Total processed documents: {len(self.docs_processed)}.")


@dataclass
class EmbeddingHandler:
    embedding_model_name: str = "sentence-transformers/all-MiniLM-L6-v2"
    device: str = "cuda" if torch.cuda.is_available() else "cpu"
    knowledge_vector_database: Optional[FAISS] = field(init=False, default=None)
    embedding_model: Optional[HuggingFaceEmbeddings] = field(init=False, default=None)

    def generate_embeddings(self, documents: List[LangchainDocument]) -> FAISS:
        """
        Generate embeddings for documents and create a FAISS index.
        """
        logger.info(f"Generating embeddings using model '{self.embedding_model_name}' on device '{self.device}'.")
        try:
            self.embedding_model = HuggingFaceEmbeddings(
                model_name=self.embedding_model_name,
                model_kwargs={"device": self.device},
                encode_kwargs={"normalize_embeddings": True}
            )
            faiss_index = FAISS.from_documents(
                documents,
                self.embedding_model,
                distance_strategy=DistanceStrategy.COSINE
            )
            logger.info("FAISS index created successfully.")
            self.knowledge_vector_database = faiss_index
            return faiss_index
        except Exception as e:
            logger.error(f"Error during embedding generation: {e}")
            raise

    def embed_query(self, query: str) -> List[float]:
        """
        Embed a user query.
        """
        if self.embedding_model is None:
            logger.error("Embedding model not initialized.")
            raise ValueError("Embedding model not initialized.")
        try:
            logger.info("Embedding user query.")
            query_vector = self.embedding_model.embed_query(query)
            logger.info("User query embedded successfully.")
            return query_vector
        except Exception as e:
            logger.error(f"Error embedding query: {e}")
            raise


@dataclass
class Visualization:
    embedding_projector: Optional[pacmap.PaCMAP] = field(init=False, default=None)
    documents_projected: Optional[np.ndarray] = field(init=False, default=None)

    def project_embeddings(self, faiss_index: FAISS, query_vector: List[float]) -> np.ndarray:
        """
        Project high-dimensional embeddings to 2D for visualization.
        """
        logger.info("Projecting embeddings to 2D space using PaCMAP.")
        try:
            embeddings = faiss_index.index.reconstruct_n(0, faiss_index.index.ntotal)
            embeddings = np.array(embeddings)
            # Append query vector
            embeddings = np.vstack([embeddings, query_vector])
            self.embedding_projector = pacmap.PaCMAP(n_components=2, random_state=42)
            self.documents_projected = self.embedding_projector.fit_transform(embeddings)
            logger.info("Projection complete.")
            return self.documents_projected
        except Exception as e:
            logger.error(f"Error during embedding projection: {e}")
            raise

    def create_dataframe(self, faiss_index: FAISS, query: str) -> pd.DataFrame:
        """
        Create a DataFrame for visualization.
        """
        logger.info("Creating DataFrame for visualization.")
        try:
            sources = [doc.metadata.get("source", "Unknown") for doc in faiss_index.index_to_docstore_id]
            extracts = [doc.page_content[:100] + "..." for doc in faiss_index.docstore.values()]
            df = pd.DataFrame({
                "x": self.documents_projected[:-1, 0],
                "y": self.documents_projected[:-1, 1],
                "source": sources,
                "extract": extracts,
                "symbol": "circle",
                "size_col": 4
            })
            # Add user query
            df = df.append({
                "x": self.documents_projected[-1, 0],
                "y": self.documents_projected[-1, 1],
                "source": "User Query",
                "extract": query,
                "symbol": "star",
                "size_col": 12
            }, ignore_index=True)
            logger.info("DataFrame created successfully.")
            return df
        except Exception as e:
            logger.error(f"Error creating DataFrame: {e}")
            raise

    def visualize(self, df: pd.DataFrame) -> None:
        """
        Visualize the embeddings using Plotly.
        """
        logger.info("Generating visualization.")
        try:
            fig = px.scatter(
                df,
                x="x",
                y="y",
                color="source",
                hover_data=["extract"],
                size="size_col",
                symbol="symbol",
                width=1000,
                height=700,
                title="2D Projection of Document Embeddings via PaCMAP"
            )
            fig.update_traces(
                marker=dict(opacity=0.7, line=dict(width=0.5, color="DarkSlateGrey")),
                selector=dict(mode="markers")
            )
            fig.update_layout(
                legend_title_text="<b>Source</b>",
                title="<b>2D Projection of Document Embeddings via PaCMAP</b>"
            )
            fig.show()
            logger.info("Visualization complete.")
        except Exception as e:
            logger.error(f"Error during visualization: {e}")
            raise


@dataclass
class RAGPipeline:
    dataset_name: str
    query: str
    split: str = "train"

    def run(self) -> None:
        """
        Run the full RAG pipeline.
        """
        try:
            # Load data
            data_loader = DataLoader(dataset_name=self.dataset_name, split=self.split)
            data_loader.load_dataset()
            text_columns = data_loader.identify_text_columns()
            if not text_columns:
                logger.error("No text columns available for processing.")
                return

            # Process documents
            doc_processor = DocumentProcessor(text_columns=text_columns)
            doc_processor.process_documents(data_loader.dataset)

            # Generate embeddings and create FAISS index
            embed_handler = EmbeddingHandler()
            faiss_index = embed_handler.generate_embeddings(doc_processor.docs_processed)

            # Embed user query
            query_vector = embed_handler.embed_query(self.query)

            # Visualization
            visualizer = Visualization()
            projection = visualizer.project_embeddings(faiss_index, query_vector)
            df = visualizer.create_dataframe(faiss_index, self.query)
            visualizer.visualize(df)

            logger.info("RAG pipeline completed successfully.")
        except Exception as e:
            logger.error(f"RAG pipeline failed: {e}")


# Example usage
if __name__ == "__main__":
    import torch

    user_query = "How to create a pipeline object?"
    dataset_name = "m-ric/huggingface_doc"  # Replace with your desired dataset

    pipeline = RAGPipeline(dataset_name=dataset_name, query=user_query)
    pipeline.run()