docanalysis.py

# Based of the sample provided https://www.ibm.com/think/tutorials/build-document-question-answering-system-with-docling-and-granite
# Required imports
import os
import tempfile
import shutil
import logging
logging.basicConfig(
         format='%(asctime)s %(levelname)-8s %(message)s',
         level=logging.INFO,
         datefmt='%Y-%m-%d %H:%M:%S')

logging.info(f"Loading IMPORTS")


from pathlib import Path
from IPython.display import Markdown, display

# LangChain imports
from langchain_community.document_loaders import UnstructuredMarkdownLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_ollama import OllamaEmbeddings, OllamaLLM
from langchain_community.vectorstores import FAISS

from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ConversationBufferMemory
from langchain_core.messages import SystemMessage

logging.info(f"Loaded IMPORTS")

def setup_qa_chain(markdown_path: Path, embeddings_model_name:str = "nomic-embed-text:latest", model_name: str = "gabegoodhart/granite3.2-preview:8b"):

    """Set up the QA chain for document processing"""
    embeddings = OllamaEmbeddings(model=embeddings_model_name)
    loader = UnstructuredMarkdownLoader(str(markdown_path))
    fname = "faiss_index"

    if (not os.path.isdir(fname)):
        logging.info(f"  Creating the vectorstore")

        logging.info(f"  Loading Document")
        documents = loader.load()
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200,
            length_function=len,
        )
        logging.info(f"  Splitting Document")
        texts = text_splitter.split_documents(documents)

        logging.info(f"  Create the vector store")


        vectorstore = FAISS.from_documents(texts, embeddings)


        vectorstore.save_local()

    else:
        logging.info(f"  Loading the vectorstore")
        vectorstore = FAISS.load_local(
            fname, embeddings, allow_dangerous_deserialization=True
        )

    logging.info(f"  Initialize LLM")

    llm = OllamaLLM(
        model=model_name,
        temperature=0
    )
    logging.info(f"  Set up conversation memory")
    memory = ConversationBufferMemory(
        memory_key="chat_history",
        output_key="answer",
        return_messages=True
    )
    logging.info(f"  Create the chainCreate the chain")
    qa_chain = ConversationalRetrievalChain.from_llm(
        llm=llm,
        retriever=vectorstore.as_retriever(search_kwargs={"k": 10}),
        memory=memory,
        return_source_documents=True
    )


    return qa_chain

def ask_question(qa_chain, question: str):
    """Ask a question and display the answer"""
    logging.info(f"Asking  .... \033[95m{question} \033[92m")
    result = qa_chain.invoke({ "question": question})
    logging.info(f"Answer  .... \033[95m{result['answer']} \033[92m")


logging.info(f"Setup QA chain  ....")
qa_chain = setup_qa_chain("doc2.md")
# Example questions
questions = [
    "How do I create a subscription object with apic cli?"
]
for question in questions:
    ask_question(qa_chain, question)