Create a RAG system using OpenVINO and LlamaIndex#

import os
import requests

os.environ["GIT_CLONE_PROTECTION_ACTIVE"] = "false"

r = requests.get(
    url="https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/notebook_utils.py",
)
with open("notebook_utils.py", "w") as f:
    f.write(r.text)

r = requests.get(
    url="https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/pip_helper.py",
)
open("pip_helper.py", "w").write(r.text)

from pip_helper import pip_install

pip_install(
    "-q",
    "--extra-index-url",
    "https://download.pytorch.org/whl/cpu",
    "llama-index",
    "faiss-cpu",
    "pymupdf",
    "langchain",
    "llama-index-readers-file",
    "llama-index-vector-stores-faiss",
    "llama-index-llms-langchain",
    "llama-index-llms-huggingface>=0.3.0,<0.3.4",
    "llama-index-embeddings-huggingface>=0.3.0",
)
pip_install("-q", "git+https://github.com/huggingface/optimum-intel.git", "git+https://github.com/openvinotoolkit/nncf.git", "datasets", "accelerate", "gradio")
pip_install("--pre", "-U", "openvino>=2024.2", "--extra-index-url", "https://storage.openvinotoolkit.org/simple/wheels/nightly")
pip_install("--pre", "-U", "openvino-tokenizers[transformers]>=2024.2", "--extra-index-url", "https://storage.openvinotoolkit.org/simple/wheels/nightly")
pip_install("-q", "--no-deps", "llama-index-llms-openvino>=0.3.1", "llama-index-embeddings-openvino>=0.2.1", "llama-index-postprocessor-openvino-rerank>=0.2.0")

import os
from pathlib import Path
import requests
import shutil
import io

# fetch model configuration

config_shared_path = Path("../../utils/llm_config.py")
config_dst_path = Path("llm_config.py")
text_example_en_path = Path("text_example_en.pdf")
text_example_cn_path = Path("text_example_cn.pdf")
text_example_en = "https://github.com/openvinotoolkit/openvino_notebooks/files/15039728/Platform.Brief_Intel.vPro.with.Intel.Core.Ultra_Final.pdf"
text_example_cn = "https://github.com/openvinotoolkit/openvino_notebooks/files/15039713/Platform.Brief_Intel.vPro.with.Intel.Core.Ultra_Final_CH.pdf"

if not config_dst_path.exists():
    if config_shared_path.exists():
        try:
            os.symlink(config_shared_path, config_dst_path)
        except Exception:
            shutil.copy(config_shared_path, config_dst_path)
    else:
        r = requests.get(url="https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/llm_config.py")
        with open("llm_config.py", "w", encoding="utf-8") as f:
            f.write(r.text)
elif not os.path.islink(config_dst_path):
    print("LLM config will be updated")
    if config_shared_path.exists():
        shutil.copy(config_shared_path, config_dst_path)
    else:
        r = requests.get(url="https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/llm_config.py")
        with open("llm_config.py", "w", encoding="utf-8") as f:
            f.write(r.text)


if not text_example_en_path.exists():
    r = requests.get(url=text_example_en)
    content = io.BytesIO(r.content)
    with open("text_example_en.pdf", "wb") as f:
        f.write(content.read())

if not text_example_cn_path.exists():
    r = requests.get(url=text_example_cn)
    content = io.BytesIO(r.content)
    with open("text_example_cn.pdf", "wb") as f:
        f.write(content.read())

from pathlib import Path
import ipywidgets as widgets
from notebook_utils import device_widget, optimize_bge_embedding

from llm_config import (
    SUPPORTED_EMBEDDING_MODELS,
    SUPPORTED_RERANK_MODELS,
    SUPPORTED_LLM_MODELS,
)

model_languages = list(SUPPORTED_LLM_MODELS)

model_language = widgets.Dropdown(
    options=model_languages,
    value=model_languages[0],
    description="Model Language:",
    disabled=False,
)

model_language

llm_model_ids = [model_id for model_id, model_config in SUPPORTED_LLM_MODELS[model_language.value].items() if model_config.get("rag_prompt_template")]

llm_model_id = widgets.Dropdown(
    options=llm_model_ids,
    value=llm_model_ids[-1],
    description="Model:",
    disabled=False,
)

llm_model_id

llm_model_configuration = SUPPORTED_LLM_MODELS[model_language.value][llm_model_id.value]
print(f"Selected LLM model {llm_model_id.value}")

from llama_index.core import VectorStoreIndex, StorageContext
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core import Settings
from llama_index.readers.file import PyMuPDFReader
from llama_index.vector_stores.faiss import FaissVectorStore
from transformers import StoppingCriteria, StoppingCriteriaList
import faiss
import torch

if model_language.value == "English":
    text_example_path = "text_example_en.pdf"
else:
    text_example_path = "text_example_cn.pdf"


class StopOnTokens(StoppingCriteria):
    def __init__(self, token_ids):
        self.token_ids = token_ids

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
        for stop_id in self.token_ids:
            if input_ids[0][-1] == stop_id:
                return True
        return False


if stop_tokens is not None:
    if isinstance(stop_tokens[0], str):
        stop_tokens = llm._tokenizer.convert_tokens_to_ids(stop_tokens)
    stop_tokens = [StopOnTokens(stop_tokens)]

loader = PyMuPDFReader()
documents = loader.load(file_path=text_example_path)

# dimensions of embedding model
d = embedding._model.request.outputs[0].get_partial_shape()[2].get_length()
faiss_index = faiss.IndexFlatL2(d)
Settings.embed_model = embedding

llm.max_new_tokens = 2048
if stop_tokens is not None:
    llm._stopping_criteria = StoppingCriteriaList(stop_tokens)
Settings.llm = llm

vector_store = FaissVectorStore(faiss_index=faiss_index)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(
    documents,
    storage_context=storage_context,
    transformations=[SentenceSplitter(chunk_size=200, chunk_overlap=40)],
)

query_engine = index.as_query_engine(streaming=True, similarity_top_k=10, node_postprocessors=[reranker])
if model_language.value == "English":
    query = "What can Intel vPro® Enterprise systems offer?"
else:
    query = "英特尔博锐® Enterprise系统提供哪些功能？"

streaming_response = query_engine.query(query)
streaming_response.print_response_stream()

prompts_dict = query_engine.get_prompts()


def display_prompt_dict(prompts_dict):
    for k, p in prompts_dict.items():
        text_md = f"**Prompt Key**: {k}<br>" f"**Text:** <br>"
        display(Markdown(text_md))
        print(p.get_template())
        display(Markdown("<br><br>"))


display_prompt_dict(prompts_dict)

from langchain.text_splitter import RecursiveCharacterTextSplitter
from llama_index.core.node_parser import LangchainNodeParser
import gradio as gr

TEXT_SPLITERS = {
    "SentenceSplitter": SentenceSplitter,
    "RecursiveCharacter": RecursiveCharacterTextSplitter,
}


def default_partial_text_processor(partial_text: str, new_text: str):
    """
    helper for updating partially generated answer, used by default

    Params:
      partial_text: text buffer for storing previosly generated text
      new_text: text update for the current step
    Returns:
      updated text string

    """
    partial_text += new_text
    return partial_text


text_processor = llm_model_configuration.get("partial_text_processor", default_partial_text_processor)


def create_vectordb(doc, spliter_name, chunk_size, chunk_overlap, vector_search_top_k, vector_rerank_top_n, run_rerank):
    """
    Initialize a vector database

    Params:
      doc: orignal documents provided by user
      chunk_size:  size of a single sentence chunk
      chunk_overlap: overlap size between 2 chunks
      vector_search_top_k: Vector search top k
      vector_rerank_top_n: Rerrank top n
      run_rerank: whether to run reranker

    """
    global query_engine
    global index

    if vector_rerank_top_n > vector_search_top_k:
        gr.Warning("Search top k must >= Rerank top n")

    loader = PyMuPDFReader()
    documents = loader.load(file_path=doc.name)
    spliter = TEXT_SPLITERS[spliter_name](chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    if spliter_name == "RecursiveCharacter":
        spliter = LangchainNodeParser(spliter)
    faiss_index = faiss.IndexFlatL2(d)
    vector_store = FaissVectorStore(faiss_index=faiss_index)
    storage_context = StorageContext.from_defaults(vector_store=vector_store)

    index = VectorStoreIndex.from_documents(
        documents,
        storage_context=storage_context,
        transformations=[spliter],
    )
    if run_rerank:
        reranker.top_n = vector_rerank_top_n
        query_engine = index.as_query_engine(streaming=True, similarity_top_k=vector_search_top_k, node_postprocessors=[reranker])
    else:
        query_engine = index.as_query_engine(streaming=True, similarity_top_k=vector_search_top_k)

    return "Vector database is Ready"


def update_retriever(vector_search_top_k, vector_rerank_top_n, run_rerank):
    """
    Update retriever

    Params:
      vector_search_top_k: size of searching results
      vector_rerank_top_n:  size of rerank results
      run_rerank: whether run rerank step

    """
    global query_engine
    global index

    if vector_rerank_top_n > vector_search_top_k:
        gr.Warning("Search top k must >= Rerank top n")

    if run_rerank:
        reranker.top_n = vector_rerank_top_n
        query_engine = index.as_query_engine(streaming=True, similarity_top_k=vector_search_top_k, node_postprocessors=[reranker])
    else:
        query_engine = index.as_query_engine(streaming=True, similarity_top_k=vector_search_top_k)


def bot(history, temperature, top_p, top_k, repetition_penalty, do_rag):
    """
    callback function for running chatbot on submit button click

    Params:
      history: conversation history
      temperature:  parameter for control the level of creativity in AI-generated text.
                    By adjusting the `temperature`, you can influence the AI model's probability distribution, making the text more focused or diverse.
      top_p: parameter for control the range of tokens considered by the AI model based on their cumulative probability.
      top_k: parameter for control the range of tokens considered by the AI model based on their cumulative probability, selecting number of tokens with highest probability.
      repetition_penalty: parameter for penalizing tokens based on how frequently they occur in the text.
      do_rag: whether do RAG when generating texts.

    """
    llm.generate_kwargs = dict(
        temperature=temperature,
        do_sample=temperature > 0.0,
        top_p=top_p,
        top_k=top_k,
        repetition_penalty=repetition_penalty,
    )

    partial_text = ""
    if do_rag:
        streaming_response = query_engine.query(history[-1][0])
        for new_text in streaming_response.response_gen:
            partial_text = text_processor(partial_text, new_text)
            history[-1][1] = partial_text
            yield history
    else:
        streaming_response = llm.stream_complete(history[-1][0])
        for new_text in streaming_response:
            partial_text = text_processor(partial_text, new_text.delta)
            history[-1][1] = partial_text
            yield history


def request_cancel():
    llm._model.request.cancel()

if not Path("gradio_helper.py").exists():
    r = requests.get(url="https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/notebooks/llm-rag-llamaindex/gradio_helper.py")
    open("gradio_helper.py", "w").write(r.text)

from gradio_helper import make_demo

demo = make_demo(
    load_doc_fn=create_vectordb,
    run_fn=bot,
    stop_fn=request_cancel,
    update_retriever_fn=update_retriever,
    model_name=llm_model_id.value,
    language=model_language.value,
)

try:
    demo.queue().launch()
except Exception:
    demo.queue().launch(share=True)
# If you are launching remotely, specify server_name and server_port
# EXAMPLE: `demo.launch(server_name='your server name', server_port='server port in int')`
# To learn more please refer to the Gradio docs: https://gradio.app/docs/

# please uncomment and run this cell for stopping gradio interface
# demo.close()