nlp-llamaindex-advanced-rag-pipeline.ipynb

{"metadata":{"kernelspec":{"display_name":"Python 3 (ipykernel)","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.10.13"}},"nbformat_minor":5,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Lesson 1: Advanced RAG Pipeline","metadata":{}},{"cell_type":"code","source":"import utils\n\nimport os\nimport openai\nopenai.api_key = utils.get_openai_api_key()","metadata":{"height":98,"tags":[]},"execution_count":32,"outputs":[]},{"cell_type":"code","source":"from llama_index import SimpleDirectoryReader\n\ndocuments = SimpleDirectoryReader(\n    input_files=[\"./eBook-How-to-Build-a-Career-in-AI.pdf\"]\n).load_data()","metadata":{"height":98,"tags":[]},"execution_count":33,"outputs":[]},{"cell_type":"code","source":"#!pip install python-dotenv\n\n\nimport os\nfrom dotenv import load_dotenv, find_dotenv\n\nimport numpy as np\nfrom trulens_eval import (\n    Feedback,\n    TruLlama,\n    OpenAI\n)\n\nfrom trulens_eval.feedback import Groundedness\nimport nest_asyncio\n\nnest_asyncio.apply()\n\n\ndef get_openai_api_key():\n    _ = load_dotenv(find_dotenv())\n\n    return os.getenv(\"OPENAI_API_KEY\")\n\n\ndef get_hf_api_key():\n    _ = load_dotenv(find_dotenv())\n\n    return os.getenv(\"HUGGINGFACE_API_KEY\")\n\nopenai = OpenAI()\n\nqa_relevance = (\n    Feedback(openai.relevance_with_cot_reasons, name=\"Answer Relevance\")\n    .on_input_output()\n)\n\nqs_relevance = (\n    Feedback(openai.relevance_with_cot_reasons, name = \"Context Relevance\")\n    .on_input()\n    .on(TruLlama.select_source_nodes().node.text)\n    .aggregate(np.mean)\n)\n\n#grounded = Groundedness(groundedness_provider=openai, summarize_provider=openai)\ngrounded = Groundedness(groundedness_provider=openai)\n\ngroundedness = (\n    Feedback(grounded.groundedness_measure_with_cot_reasons, name=\"Groundedness\")\n        .on(TruLlama.select_source_nodes().node.text)\n        .on_output()\n        .aggregate(grounded.grounded_statements_aggregator)\n)\n\nfeedbacks = [qa_relevance, qs_relevance, groundedness]\n\ndef get_trulens_recorder(query_engine, feedbacks, app_id):\n    tru_recorder = TruLlama(\n        query_engine,\n        app_id=app_id,\n        feedbacks=feedbacks\n    )\n    return tru_recorder\n\ndef get_prebuilt_trulens_recorder(query_engine, app_id):\n    tru_recorder = TruLlama(\n        query_engine,\n        app_id=app_id,\n        feedbacks=feedbacks\n        )\n    return tru_recorder\n\nfrom llama_index import ServiceContext, VectorStoreIndex, StorageContext\nfrom llama_index.node_parser import SentenceWindowNodeParser\nfrom llama_index.indices.postprocessor import MetadataReplacementPostProcessor\nfrom llama_index.indices.postprocessor import SentenceTransformerRerank\nfrom llama_index import load_index_from_storage\nimport os\n\n\ndef build_sentence_window_index(\n    document, llm, embed_model=\"local:BAAI/bge-small-en-v1.5\", save_dir=\"sentence_index\"\n):\n    # create the sentence window node parser w/ default settings\n    node_parser = SentenceWindowNodeParser.from_defaults(\n        window_size=3,\n        window_metadata_key=\"window\",\n        original_text_metadata_key=\"original_text\",\n    )\n    sentence_context = ServiceContext.from_defaults(\n        llm=llm,\n        embed_model=embed_model,\n        node_parser=node_parser,\n    )\n    if not os.path.exists(save_dir):\n        sentence_index = VectorStoreIndex.from_documents(\n            [document], service_context=sentence_context\n        )\n        sentence_index.storage_context.persist(persist_dir=save_dir)\n    else:\n        sentence_index = load_index_from_storage(\n            StorageContext.from_defaults(persist_dir=save_dir),\n            service_context=sentence_context,\n        )\n\n    return sentence_index\n\n\ndef get_sentence_window_query_engine(\n    sentence_index,\n    similarity_top_k=6,\n    rerank_top_n=2,\n):\n    # define postprocessors\n    postproc = MetadataReplacementPostProcessor(target_metadata_key=\"window\")\n    rerank = SentenceTransformerRerank(\n        top_n=rerank_top_n, model=\"BAAI/bge-reranker-base\"\n    )\n\n    sentence_window_engine = sentence_index.as_query_engine(\n        similarity_top_k=similarity_top_k, node_postprocessors=[postproc, rerank]\n    )\n    return sentence_window_engine\n\n\nfrom llama_index.node_parser import HierarchicalNodeParser\n\nfrom llama_index.node_parser import get_leaf_nodes\nfrom llama_index import StorageContext\nfrom llama_index.retrievers import AutoMergingRetriever\nfrom llama_index.indices.postprocessor import SentenceTransformerRerank\nfrom llama_index.query_engine import RetrieverQueryEngine\n\n\ndef build_automerging_index(\n    documents,\n    llm,\n    embed_model=\"local:BAAI/bge-small-en-v1.5\",\n    save_dir=\"merging_index\",\n    chunk_sizes=None,\n):\n    chunk_sizes = chunk_sizes or [2048, 512, 128]\n    node_parser = HierarchicalNodeParser.from_defaults(chunk_sizes=chunk_sizes)\n    nodes = node_parser.get_nodes_from_documents(documents)\n    leaf_nodes = get_leaf_nodes(nodes)\n    merging_context = ServiceContext.from_defaults(\n        llm=llm,\n        embed_model=embed_model,\n    )\n    storage_context = StorageContext.from_defaults()\n    storage_context.docstore.add_documents(nodes)\n\n    if not os.path.exists(save_dir):\n        automerging_index = VectorStoreIndex(\n            leaf_nodes, storage_context=storage_context, service_context=merging_context\n        )\n        automerging_index.storage_context.persist(persist_dir=save_dir)\n    else:\n        automerging_index = load_index_from_storage(\n            StorageContext.from_defaults(persist_dir=save_dir),\n            service_context=merging_context,\n        )\n    return automerging_index\n\n\ndef get_automerging_query_engine(\n    automerging_index,\n    similarity_top_k=12,\n    rerank_top_n=2,\n):\n    base_retriever = automerging_index.as_retriever(similarity_top_k=similarity_top_k)\n    retriever = AutoMergingRetriever(\n        base_retriever, automerging_index.storage_context, verbose=True\n    )\n    rerank = SentenceTransformerRerank(\n        top_n=rerank_top_n, model=\"BAAI/bge-reranker-base\"\n    )\n    auto_merging_engine = RetrieverQueryEngine.from_args(\n        retriever, node_postprocessors=[rerank]\n    )\n    return auto_merging_engine\n","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(type(documents), \"\\n\")\nprint(len(documents), \"\\n\")\nprint(type(documents[0]))\nprint(documents[0])","metadata":{"height":81,"tags":[]},"execution_count":34,"outputs":[{"name":"stdout","output_type":"stream","text":"<class 'list'> \n\n\n\n41 \n\n\n\n<class 'llama_index.schema.Document'>\n\nDoc ID: 7af5e50d-5a68-4811-b7f6-fed7c09df94b\n\nText: PAGE 1Founder, DeepLearning.AICollected Insights from Andrew Ng\n\nHow to  Build Your Career in AIA Simple Guide\n"}]},{"cell_type":"markdown","source":"## Basic RAG pipeline","metadata":{}},{"cell_type":"code","source":"from llama_index import Document\n\ndocument = Document(text=\"\\n\\n\".join([doc.text for doc in documents]))","metadata":{"height":64,"tags":[]},"execution_count":35,"outputs":[]},{"cell_type":"code","source":"from llama_index import VectorStoreIndex\nfrom llama_index import ServiceContext\nfrom llama_index.llms import OpenAI\n\nllm = OpenAI(model=\"gpt-3.5-turbo\", temperature=0.1)\nservice_context = ServiceContext.from_defaults(\n    llm=llm, embed_model=\"local:BAAI/bge-small-en-v1.5\"\n)\nindex = VectorStoreIndex.from_documents([document],\n                                        service_context=service_context)","metadata":{"height":183,"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"query_engine = index.as_query_engine()","metadata":{"height":30,"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"response = query_engine.query(\n    \"What are steps to take when finding projects to build your experience?\"\n)\nprint(str(response))","metadata":{"height":81,"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Evaluation setup using TruLens","metadata":{}},{"cell_type":"code","source":"eval_questions = []\nwith open('eval_questions.txt', 'r') as file:\n    for line in file:\n        # Remove newline character and convert to integer\n        item = line.strip()\n        print(item)\n        eval_questions.append(item)","metadata":{"height":132,"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# You can try your own question:\nnew_question = \"What is the right AI job for me?\"\neval_questions.append(new_question)","metadata":{"height":64},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(eval_questions)","metadata":{"height":30},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from trulens_eval import Tru\ntru = Tru()\n\ntru.reset_database()","metadata":{"height":81,"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"For the classroom, we've written some of the code in helper functions inside a utils.py file.  \n- You can view the utils.py file in the file directory by clicking on the \"Jupyter\" logo at the top of the notebook.\n- In later lessons, you'll get to work directly with the code that's currently wrapped inside these helper functions, to give you more options to customize your RAG pipeline.","metadata":{}},{"cell_type":"code","source":"from utils import get_prebuilt_trulens_recorder\n\ntru_recorder = get_prebuilt_trulens_recorder(query_engine,\n                                             app_id=\"Direct Query Engine\")","metadata":{"height":81,"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with tru_recorder as recording:\n    for question in eval_questions:\n        response = query_engine.query(question)","metadata":{"height":64,"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"records, feedback = tru.get_records_and_feedback(app_ids=[])","metadata":{"height":30,"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"records.head()","metadata":{"height":30,"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# launches on http://localhost:8501/\ntru.run_dashboard()","metadata":{"height":47,"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Advanced RAG pipeline","metadata":{}},{"cell_type":"markdown","source":"### 1. Sentence Window retrieval","metadata":{}},{"cell_type":"code","source":"from llama_index.llms import OpenAI\n\nllm = OpenAI(model=\"gpt-3.5-turbo\", temperature=0.1)","metadata":{"height":64,"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from utils import build_sentence_window_index\n\nsentence_index = build_sentence_window_index(\n    document,\n    llm,\n    embed_model=\"local:BAAI/bge-small-en-v1.5\",\n    save_dir=\"sentence_index\"\n)","metadata":{"height":149,"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from utils import get_sentence_window_query_engine\n\nsentence_window_engine = get_sentence_window_query_engine(sentence_index)","metadata":{"height":64,"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"window_response = sentence_window_engine.query(\n    \"how do I get started on a personal project in AI?\"\n)\nprint(str(window_response))","metadata":{"height":81,"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"tru.reset_database()\n\ntru_recorder_sentence_window = get_prebuilt_trulens_recorder(\n    sentence_window_engine,\n    app_id = \"Sentence Window Query Engine\"\n)","metadata":{"height":115,"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for question in eval_questions:\n    with tru_recorder_sentence_window as recording:\n        response = sentence_window_engine.query(question)\n        print(question)\n        print(str(response))","metadata":{"height":98,"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"tru.get_leaderboard(app_ids=[])","metadata":{"height":30,"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# launches on http://localhost:8501/\ntru.run_dashboard()","metadata":{"height":47,"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 2. Auto-merging retrieval","metadata":{}},{"cell_type":"code","source":"from utils import build_automerging_index\n\nautomerging_index = build_automerging_index(\n    documents,\n    llm,\n    embed_model=\"local:BAAI/bge-small-en-v1.5\",\n    save_dir=\"merging_index\"\n)","metadata":{"height":149,"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from utils import get_automerging_query_engine\n\nautomerging_query_engine = get_automerging_query_engine(\n    automerging_index,\n)","metadata":{"height":98,"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"auto_merging_response = automerging_query_engine.query(\n    \"How do I build a portfolio of AI projects?\"\n)\nprint(str(auto_merging_response))","metadata":{"height":81,"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"tru.reset_database()\n\ntru_recorder_automerging = get_prebuilt_trulens_recorder(automerging_query_engine,\n                                                         app_id=\"Automerging Query Engine\")","metadata":{"height":81,"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for question in eval_questions:\n    with tru_recorder_automerging as recording:\n        response = automerging_query_engine.query(question)\n        print(question)\n        print(response)","metadata":{"height":98,"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"tru.get_leaderboard(app_ids=[])","metadata":{"height":30,"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# launches on http://localhost:8501/\ntru.run_dashboard()","metadata":{"height":47,"tags":[]},"execution_count":null,"outputs":[]}]}