OpenVINO と LlamaIndex を使用して RAG システムを作成#

import os 

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

%pip install -Uq pip 
%pip uninstall -q -y optimum optimum-intel 
%pip install -q --extra-index-url https://download.pytorch.org/whl/cpu\ 
"llama-index" "faiss-cpu" "pymupdf" "llama-index-readers-file" "llama-index-vector-stores-faiss" "llama-index-llms-langchain" "llama-index-llms-openvino" "llama-index-embeddings-openvino" "llama-index-postprocessor-openvino-rerank" "transformers>=4.40" \ 
"git+https://github.com/huggingface/optimum-intel.git"\ 
"git+https://github.com/openvinotoolkit/nncf.git"\ 
"datasets"\ 
"accelerate"\ 
"gradio" \ 
"langchain" 
%pip install --pre -Uq openvino openvino-tokenizers[transformers] --extra-index-url https://storage.openvinotoolkit.org/simple/wheels/nightly

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

# モデル構成を取得 

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 openvino as ov 
import ipywidgets as widgets

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" 
stop_tokens = llm_model_configuration.get("stop_tokens") 

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) 

# 埋め込みモデルの次元 
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, 
} 

chinese_examples = [ 
    ["英特尔®酷睿™ Ultra处理器可以降低多少功耗？"], 
    ["相比英特尔之前的移动处理器产品，英特尔®酷睿™ Ultra处理器的AI推理性能提升了多少？"], 
    ["英特尔博锐® Enterprise系统提供哪些功能？"], 
] 

english_examples = [ 
    ["How much power consumption can Intel® Core™ Ultra Processors help save?"], 
    ["Compared to Intel’s previous mobile processor, what is the advantage of Intel® Core™ Ultra Processors for Artificial Intelligence?"], 
    ["What can Intel vPro® Enterprise systems offer?"], 
] 

examples = chinese_examples if (model_language.value == "Chinese") else english_examples 

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 user(message, history): 
    """ 
    callback function for updating user messages in interface on submit button click 

    Params: 
        message: current message 
        history: conversation history 
    Returns:
        None 
    """ 
    # ユーザーのメッセージを会話履歴に追加 
    return "", history + [[message, ""]] 

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() 

def clear_files(): 
    return "Vector Store is Not ready" 

with gr.Blocks( 
    theme=gr.themes.Soft(), 
    css=".disclaimer {font-variant-caps: all-small-caps;}", 
) as demo: 
    gr.Markdown("""<h1><center>QA over Document</center></h1>""") 
    gr.Markdown(f"""<center>Powered by OpenVINO and {llm_model_id.value} </center>""") 
    with gr.Row(): 
        with gr.Column(scale=1): 
            docs = gr.File( 
                label="Step 1: Load a PDF file", 
                value=text_example_path, 
                file_types=[ 
                    ".pdf", 
                ], 
            ) 
            load_docs = gr.Button("Step 2: Build Vector Store", variant="primary") 
            db_argument = gr.Accordion("Vector Store Configuration", open=False) 
            with db_argument: 
                spliter = gr.Dropdown( 
                    ["SentenceSplitter", "RecursiveCharacter"], 
                    value="SentenceSplitter", 
                    label="Text Spliter", 
                    info="Method used to splite the documents", 
                    multiselect=False, 
                ) 

                chunk_size = gr.Slider( 
                    label="Chunk size", 
                    value=200, 
                    minimum=50, 
                    maximum=2000, 
                    step=50, 
                    interactive=True, 
                    info="Size of sentence chunk", 
                ) 

                chunk_overlap = gr.Slider( 
                    label="Chunk overlap", 
                    value=20, 
                    minimum=0, 
                    maximum=400, 
                    step=10, 
                    interactive=True, 
                    info=("Overlap between 2 chunks"), 
                ) 

            vector_store_status = gr.Textbox( 
                label="Vector Store Status", 
                value="Vector Store is Ready", 
                interactive=False, 
            ) 
            do_rag = gr.Checkbox( 
                value=True, 
                label="RAG is ON", 
                interactive=True, 
                info="Whether to do RAG for generation", 
            ) 
            with gr.Accordion("Generation Configuration", open=False): 
                with gr.Row(): 
                    with gr.Column(): 
                        with gr.Row(): 
                            temperature = gr.Slider( 
                                label="Temperature", 
                                value=0.1, 
                                minimum=0.0, 
                                maximum=1.0, 
                                step=0.1, 
                                interactive=True, 
                                info="Higher values produce more diverse outputs", 
                            ) 
                    with gr.Column(): 
                        with gr.Row(): 
                            top_p = gr.Slider( 
                                label="Top-p (nucleus sampling)", 
                                value=1.0, 
                                minimum=0.0, 
                                maximum=1, 
                                step=0.01, 
                                interactive=True, 
                                info=( 
                                    "Sample from the smallest possible set of tokens whose cumulative probability " 
                                    "exceeds top_p.Set to 1 to disable and sample from all tokens."                                 ), 
                            ) 
                with gr.Column(): 
                    with gr.Row(): 
                        top_k = gr.Slider( 
                            label="Top-k", 
                            value=50, 
                            minimum=0.0, 
                            maximum=200, 
                            step=1, 
                            interactive=True, 
                            info="Sample from a shortlist of top-k tokens — 0 to disable and sample from all tokens.", 
                        ) 
                with gr.Column(): 
                    with gr.Row(): 
                        repetition_penalty = gr.Slider( 
                            label="Repetition Penalty", 
                            value=1.1, 
                            minimum=1.0, 
                            maximum=2.0, 
                            step=0.1, 
                            interactive=True, 
                            info="Penalize repetition — 1.0 to disable.", 
                        ) 
        with gr.Column(scale=4): 
            chatbot = gr.Chatbot( 
                height=600, 
                label="Step 3: Input Query", 
            ) 
            with gr.Row(): 
                with gr.Column(): 
                    with gr.Row(): 
                        msg = gr.Textbox( 
                            label="QA Message Box", 
                            placeholder="Chat Message Box", 
                            show_label=False, 
                            container=False, 
                        ) 
                with gr.Column(): 
                    with gr.Row(): 
                        submit = gr.Button("Submit", variant="primary") 
                        stop = gr.Button("Stop") 
                        clear = gr.Button("Clear") 
        gr.Examples(examples, inputs=msg, label="Click on any example and press the 'Submit' button") 
        retriever_argument = gr.Accordion("Retriever Configuration", open=True) 
        with retriever_argument: with gr.Row(): 
            with gr.Row(): 
                do_rerank = gr.Checkbox( 
                    value=True, 
                    label="Rerank searching result", 
                    interactive=True, 
                ) 
            with gr.Row(): 
                vector_rerank_top_n = gr.Slider( 
                    1, 
                    10, 
                    value=2, 
                    step=1, 
                    label="Rerank top n", 
                    info="Number of rerank results", 
                    interactive=True, 
                ) 
            with gr.Row(): 
                vector_search_top_k = gr.Slider( 
                    1, 
                    50, 
                    value=10, 
                    step=1, 
                    label="Search top k", 
                    info="Search top k must >= Rerank top n", 
                    interactive=True, 
                ) 
    docs.clear(clear_files, outputs=[vector_store_status], queue=False) 
    load_docs.click( 
        create_vectordb, 
        inputs=[docs, spliter, chunk_size, chunk_overlap, vector_search_top_k, vector_rerank_top_n, do_rerank], 
        outputs=[vector_store_status], 
        queue=False, 
    ) 
    submit_event = msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then( 
        bot, 
        [chatbot, temperature, top_p, top_k, repetition_penalty, do_rag], 
        chatbot, 
        queue=True, 
    )
    submit_click_event = submit.click(user, [msg, chatbot], [msg, chatbot], queue=False).then( 
        bot, 
        [chatbot, temperature, top_p, top_k, repetition_penalty, do_rag], 
        chatbot, 
        queue=True, 
    ) 
    stop.click( 
        fn=request_cancel, 
        inputs=None, 
        outputs=None, 
        cancels=[submit_event, submit_click_event], 
        queue=False, 
    ) 
    clear.click(lambda: None, None, chatbot, queue=False) 
    vector_search_top_k.release( 
        update_retriever, 
        [vector_search_top_k, vector_rerank_top_n, do_rerank], 
    ) 
    vector_rerank_top_n.release( 
        update_retriever, 
        [vector_search_top_k, vector_rerank_top_n, do_rerank], 
    ) 
    do_rerank.change( 
        update_retriever, 
        [vector_search_top_k, vector_rerank_top_n, do_rerank], 
    ) 

demo.queue() 
# リモートで起動する場合は、server_name と server_port を指定 
# demo.launch(server_name='your server name', server_port='server port in int') 
# プラットフォーム上で起動する際に問題がある場合は、起動メソッドに share=True を渡すことができます: 
# demo.launch(share=True) 
# インターフェイスの公開共有可能なリンクを作成。詳細はドキュメントをご覧ください: https://gradio.app/docs/ 
demo.launch()

# gradio インターフェイスを停止するにはこのセルを実行 
demo.close()