HTML & latex to Markdown

code/modules/config/constants.py

@@ -6,6 +6,8 @@ load_dotenv()
 # API Keys - Loaded from the .env file
 
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+GPT4o_API_KEY = os.getenv("GPT4o_API_KEY")
+LLAMA_CLOUD_API_KEY = os.getenv("LLAMA_CLOUD_API_KEY")
 HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
 LITERAL_API_KEY = os.getenv("LITERAL_API_KEY")
 
@@ -14,7 +16,8 @@ opening_message = f"Hey, What Can I Help You With?\n\nYou can me ask me question
 # Prompt Templates
 
 openai_prompt_template = """Use the following pieces of information to answer the user's question.
-If you don't know the answer, just say that you don't know.
+You are an intelligent chatbot designed to help students with questions regarding the course. Render math equations in LaTeX format between $$ signs, and explain the parameters and variables in the equations.
+If you don't know the answer, just say that you don't know. 
 
 Context: {context}
 Question: {question}
@@ -24,7 +27,10 @@ Helpful answer:
 """
 
 openai_prompt_template_with_history = """Use the following pieces of information to answer the user's question.
+You are an intelligent chatbot designed to help students with questions regarding the course. Render math equations in LaTeX format between $$ signs.
+
 If you don't know the answer, just say that you don't know, don't try to make up an answer.
+
 Use the history to answer the question if you can.
 Chat History:
 {chat_history}
@@ -37,7 +43,7 @@ Helpful answer:
 
 tinyllama_prompt_template = """
 <|im_start|>system
-Assistant is an intelligent chatbot designed to help students with questions regarding the course. Only answer questions using the context below and if you're not sure of an answer, you can say "I don't know". Always give a breif and concise answer to the question. Use the history to answer the question if you can.
+Assistant is an intelligent chatbot designed to help students with questions regarding the course. Only answer questions using the context below and if you're not sure of an answer, you can say "I don't know". Always give a brief and concise answer to the question. When asked for formulas, give a brief description of the formula and output math equations in LaTeX format between $ signs.
 
 Context:
 {context}
@@ -56,7 +62,7 @@ Question: {question}
 
 tinyllama_prompt_template_with_history = """
 <|im_start|>system
-Assistant is an intelligent chatbot designed to help students with questions regarding the course. Only answer questions using the context below and if you're not sure of an answer, you can say "I don't know". Always give a breif and concise answer to the question.
+Assistant is an intelligent chatbot designed to help students with questions regarding the course. Only answer questions using the context below and if you're not sure of an answer, you can say "I don't know". Always give a brief and concise answer to the question. Output math equations in LaTeX format between $ signs. Use the history to answer the question if you can.
 
 Chat History:
 {chat_history}

code/modules/data_loader.py

@@ -0,0 +1,392 @@
+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 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 $$."
+        )
+
+    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)