code/modules/data_loader.py

import os
import bs4
from urllib.parse import urljoin
import requests
import pysrt
from langchain_community.document_loaders import (
    PyMuPDFLoader,
    Docx2txtLoader,
    YoutubeLoader,
    WebBaseLoader,
    TextLoader,
)
import html2text
from langchain_community.document_loaders import UnstructuredMarkdownLoader
from llama_parse import LlamaParse
from langchain.schema import Document
import logging
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_experimental.text_splitter import SemanticChunker
from langchain_openai.embeddings import OpenAIEmbeddings
from ragatouille import RAGPretrainedModel
from langchain.chains import LLMChain
from langchain.llms import OpenAI
from langchain import PromptTemplate

try:
    from modules.helpers import get_lecture_metadata
    from modules.constants import OPENAI_API_KEY, LLAMA_CLOUD_API_KEY
except:
    from helpers import get_lecture_metadata
    from constants import OPENAI_API_KEY, LLAMA_CLOUD_API_KEY

logger = logging.getLogger(__name__)


class PDFReader:
    def __init__(self):
        pass

    def get_loader(self, pdf_path):
        loader = PyMuPDFLoader(pdf_path)
        return loader

    def get_documents(self, loader):
        return loader.load()


class LlamaParser:
    def __init__(self):
        self.parser = LlamaParse(
            api_key=LLAMA_CLOUD_API_KEY,
            result_type="markdown",
            verbose=True,
            language="en",
            gpt4o_mode=True,
            gpt4o_api_key=OPENAI_API_KEY,
            parsing_instruction="The provided documents are PDFs of lecture slides of deep learning material. They contain LaTeX equations, images, and text. The goal is to extract the text, images and equations from the slides and convert them to markdown format. The markdown should be clean and easy to read, and any math equation should be converted to LaTeX, between $$. For images, give a description and if you can, a source."
        )

    def parse(self, pdf_path):
        documents = self.parser.load_data(pdf_path)
        documents = [document.to_langchain_format() for document in documents]
        return documents


class HTMLReader:
    def __init__(self):
        pass

    def read_url(self, url):
        response = requests.get(url)
        if response.status_code == 200:
            return response.text
        else:
            logger.warning(f"Failed to download HTML from URL: {url}")
            return None

    def check_links(self, base_url, html_content):
        soup = bs4.BeautifulSoup(html_content, "html.parser")
        for link in soup.find_all("a"):
            href = link.get("href")

            if not href or href.startswith("#"):
                continue
            elif not href.startswith("https"):
                href = href.replace("http", "https")

            absolute_url = urljoin(base_url, href)
            link['href'] = absolute_url

            resp = requests.head(absolute_url)
            if resp.status_code != 200:
                logger.warning(f"Link {absolute_url} is broken")
                logger.warning(f"Status code: {resp.status_code}")

        return str(soup)

    def html_to_md(self, url, html_content):
        html_processed = self.check_links(url, html_content)
        markdown_content = html2text.html2text(html_processed)
        return markdown_content

    def read_html(self, url):
        html_content = self.read_url(url)
        if html_content:
            return self.html_to_md(url, html_content)
        else:
            return None


class FileReader:
    def __init__(self, kind):
        self.kind = kind
        if kind == "llama":
            self.pdf_reader = LlamaParser()
        else:
            self.pdf_reader = PDFReader()
        self.web_reader = HTMLReader()

    def extract_text_from_pdf(self, pdf_path):
        text = ""
        with open(pdf_path, "rb") as file:
            reader = PyPDF2.PdfReader(file)
            num_pages = len(reader.pages)
            for page_num in range(num_pages):
                page = reader.pages[page_num]
                text += page.extract_text()
        return text

    def download_pdf_from_url(self, pdf_url):
        response = requests.get(pdf_url)
        if response.status_code == 200:
            with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file:
                temp_file.write(response.content)
                temp_file_path = temp_file.name
            return temp_file_path
        else:
            print("Failed to download PDF from URL:", pdf_url)
            return None

    def read_pdf(self, temp_file_path: str):
        if self.kind == "llama":
            documents = self.pdf_reader.parse(temp_file_path)
        else:
            loader = self.pdf_reader.get_loader(temp_file_path)
            documents = self.pdf_reader.get_documents(loader)
        return documents

    def read_txt(self, temp_file_path: str):
        loader = TextLoader(temp_file_path, autodetect_encoding=True)
        return loader.load()

    def read_docx(self, temp_file_path: str):
        loader = Docx2txtLoader(temp_file_path)
        return loader.load()

    def read_srt(self, temp_file_path: str):
        subs = pysrt.open(temp_file_path)
        text = ""
        for sub in subs:
            text += sub.text
        return [Document(page_content=text)]

    def read_youtube_transcript(self, url: str):
        loader = YoutubeLoader.from_youtube_url(
            url, add_video_info=True, language=["en"], translation="en"
        )
        return loader.load()

    def read_html(self, url: str):
        return [Document(page_content=self.web_reader.read_html(url))]


class ChunkProcessor:
    def __init__(self, config):
        self.config = config

        if config["splitter_options"]["use_splitter"]:
            if config["splitter_options"]["split_by_token"]:
                self.splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
                    chunk_size=config["splitter_options"]["chunk_size"],
                    chunk_overlap=config["splitter_options"]["chunk_overlap"],
                    separators=config["splitter_options"]["chunk_separators"],
                    disallowed_special=(),
                )
            else:
                self.splitter = RecursiveCharacterTextSplitter(
                    chunk_size=config["splitter_options"]["chunk_size"],
                    chunk_overlap=config["splitter_options"]["chunk_overlap"],
                    separators=config["splitter_options"]["chunk_separators"],
                    disallowed_special=(),
                )
        else:
            self.splitter = None
        logger.info("ChunkProcessor instance created")

    def remove_delimiters(self, document_chunks: list):
        for chunk in document_chunks:
            for delimiter in self.config["splitter_options"]["delimiters_to_remove"]:
                chunk.page_content = re.sub(delimiter, " ", chunk.page_content)
        return document_chunks

    def remove_chunks(self, document_chunks: list):
        front = self.config["splitter_options"]["front_chunk_to_remove"]
        end = self.config["splitter_options"]["last_chunks_to_remove"]
        for _ in range(front):
            del document_chunks[0]
        for _ in range(end):
            document_chunks.pop()
        logger.info(f"\tNumber of pages after skipping: {len(document_chunks)}")
        return document_chunks

    def process_chunks(
            self, documents, file_type="txt", source="", page=0, metadata={}
    ):
        documents = [Document(page_content=documents, source=source, page=page)]
        if file_type == "txt":
            document_chunks = self.splitter.split_documents(documents)
        elif file_type == "pdf":
            document_chunks = documents  # Full page for now

        # add the source and page number back to the metadata
        for chunk in document_chunks:
            chunk.metadata["source"] = source
            chunk.metadata["page"] = page

            # add the metadata extracted from the document
            for key, value in metadata.items():
                chunk.metadata[key] = value

        if self.config["splitter_options"]["remove_leftover_delimiters"]:
            document_chunks = self.remove_delimiters(document_chunks)
        if self.config["splitter_options"]["remove_chunks"]:
            document_chunks = self.remove_chunks(document_chunks)

        return document_chunks

    def get_chunks(self, file_reader, uploaded_files, weblinks):
        self.document_chunks_full = []
        self.parent_document_names = []
        self.child_document_names = []
        self.documents = []
        self.document_metadata = []

        lecture_metadata = get_lecture_metadata(
            "https://dl4ds.github.io/sp2024/lectures/",
            "https://dl4ds.github.io/sp2024/schedule/",
        )  # TODO: Use more efficiently

        for file_index, file_path in enumerate(uploaded_files):
            file_name = os.path.basename(file_path)
            file_type = file_name.split(".")[-1].lower()

            # try:
            if file_type == "pdf":
                documents = file_reader.read_pdf(file_path)
            elif file_type == "txt":
                documents = file_reader.read_txt(file_path)
            elif file_type == "docx":
                documents = file_reader.read_docx(file_path)
            elif file_type == "srt":
                documents = file_reader.read_srt(file_path)
            else:
                logger.warning(f"Unsupported file type: {file_type}")
                continue

            # full_text = ""
            # for doc in documents:
            #     full_text += doc.page_content
            #     break  # getting only first page for now

            # extracted_metadata = self.extract_metadata(full_text)

            for doc in documents:
                page_num = doc.metadata.get("page", 0)
                self.documents.append(doc.page_content)
                self.document_metadata.append({"source": file_path, "page": page_num})
                if "lecture" in file_path.lower():
                    metadata = lecture_metadata.get(file_path, {})
                    metadata["source_type"] = "lecture"
                    self.document_metadata[-1].update(metadata)
                else:
                    metadata = {"source_type": "other"}

                self.child_document_names.append(f"{file_name}_{page_num}")

                self.parent_document_names.append(file_name)
                if self.config["embedding_options"]["db_option"] not in ["RAGatouille"]:
                    document_chunks = self.process_chunks(
                        self.documents[-1],
                        file_type,
                        source=file_path,
                        page=page_num,
                        metadata=metadata,
                    )
                    self.document_chunks_full.extend(document_chunks)

            # except Exception as e:
            #     logger.error(f"Error processing file {file_name}: {str(e)}")

        self.process_weblinks(file_reader, weblinks)

        logger.info(
            f"Total document chunks extracted: {len(self.document_chunks_full)}"
        )
        return (
            self.document_chunks_full,
            self.child_document_names,
            self.documents,
            self.document_metadata,
        )

    def process_weblinks(self, file_reader, weblinks):
        if weblinks[0] != "":
            logger.info(f"Splitting weblinks: total of {len(weblinks)}")

            for link_index, link in enumerate(weblinks):
                try:
                    logger.info(f"\tSplitting link {link_index + 1} : {link}")
                    if "youtube" in link:
                        documents = file_reader.read_youtube_transcript(link)
                    else:
                        documents = file_reader.read_html(link)
                        print(f"Link: {link}")
                        print(documents)
                    for doc in documents:
                        page_num = doc.metadata.get("page", 0)
                        self.documents.append(doc.page_content)
                        self.document_metadata.append(
                            {"source": link, "page": page_num}
                        )
                        self.child_document_names.append(f"{link}")

                    self.parent_document_names.append(link)
                    if self.config["embedding_options"]["db_option"] not in [
                        "RAGatouille"
                    ]:
                        document_chunks = self.process_chunks(
                            self.documents[-1],
                            "txt",
                            source=link,
                            page=0,
                            metadata={"source_type": "webpage"},
                        )
                        self.document_chunks_full.extend(document_chunks)
                except Exception as e:
                    logger.error(
                        f"Error splitting link {link_index + 1} : {link}: {str(e)}"
                    )


class DataLoader:
    def __init__(self, config):
        if config["llm_params"]["pdf_reader"] == "llama":
            if LLAMA_CLOUD_API_KEY == None or OPENAI_API_KEY == None:
                raise ValueError(
                    "Please set the LLAMA_CLOUD_API_KEY and GPT4o_API_KEY environment variables"
                )

        self.file_reader = FileReader(kind=config["llm_params"]["pdf_reader"])
        self.chunk_processor = ChunkProcessor(config)

    def get_chunks(self, uploaded_files, weblinks):
        return self.chunk_processor.get_chunks(
            self.file_reader, uploaded_files, weblinks
        )


if __name__ == "__main__":
    # read config.yml file
    import yaml
    import os
    BASE_DIR = os.path.dirname(os.path.abspath(__file__))

    with open(os.path.join(BASE_DIR, "../", "config.yml"), "r") as f:
        config = yaml.safe_load(f)

    # create DataLoader instance
    chunk_processor = ChunkProcessor(config)
    file_reader = FileReader(kind=config["llm_params"]["pdf_reader"])

    weblinks = ["https://dl4ds.github.io/sp2024/"]

    uploaded_files = []

    # get document chunks
    document_chunks, child_document_names, documents, document_metadata = chunk_processor.get_chunks(
        file_reader, uploaded_files, weblinks
    )


    print(document_chunks)