Databricks Rerankerを試してみた【RAG】

%pip install langchain databricks-vectorsearch "mlflow[databricks]>=3.4.0"

dbutils.library.restartPython()

import os

import time

import glob

from databricks.sdk import WorkspaceClient

volume_path = "/Volumes/workspace/default/rag_eval_dataset/"

w = WorkspaceClient()

client = w.serving_endpoints.get_open_ai_client()

# LLMモデルを指定

model_name = "databricks-meta-llama-3-3-70b-instruct"

# プロンプトで作成するドキュメントの内容を指示

prompts = [

"Please create a detailed internal document of about 10000 characters regarding the work regulations of 〇△× Corporation.",

"Please create a long internal document of about 10000 characters regarding the information security policy of 〇△× Corporation.",

~~~ 省略 ~~~

]

for i, prompt in enumerate(prompts, 1):

response = client.chat.completions.create(

model=model_name,

messages=[

{"role": "system", "content": "You are an HR representative at a major Japanese company."},

{"role": "user", "content": prompt},

],

)

# 1回目の生成結果を取得

initial_text = response.choices[0].message.content

# さらに文章を拡張するようにLLMに指示

expand_prompt = (

"Please expand and enrich the following internal document, adding more details, examples, and explanations to make it even longer and more comprehensive:\n"

+ initial_text

)

expand_response = client.chat.completions.create(

model=model_name,

messages=[

{"role": "system", "content": "You are an HR representative at a major Japanese company."},

{"role": "user", "content": expand_prompt},

],

)

text = expand_response.choices[0].message.content

    # ドキュメントの内容からタイトルを作成するように指示

title_prompt = "Please provide a concise and appropriate title (no more than 50 characters) for the following internal document:\n" + text

title_response = client.chat.completions.create(

model=model_name,

messages=[

{"role": "system", "content": "You are an HR representative at a major Japanese company."},

{"role": "user", "content": title_prompt},

],

)

title = title_response.choices[0].message.content.strip().replace(" ", "_").replace("/", "_")

# ドキュメントの保存

file_path = os.path.join(volume_path, f"{title}_internal_doc_{i}.txt")

with open(file_path, "w", encoding="utf-8") as f:

f.write(text)

time.sleep(1)

# ドキュメントの読み込み

def load_text_files(folder_path, extension="txt"):

docs = []

for filepath in glob.glob(os.path.join(folder_path, f"*.{extension}")):

with open(filepath, "r", encoding="utf-8") as f:

docs.append({"filename": os.path.basename(filepath), "content": f.read()})

return docs

docs = load_text_files(volume_path)

from pyspark.sql.types import StructType, StructField, StringType

from datetime import datetime, timedelta

import uuid

from langchain.text_splitter import CharacterTextSplitter

schema = StructType([

StructField("id", StringType(), False),

StructField("content", StringType(), False),

StructField("title", StringType(), True),

StructField("updated_at", StringType(), True),

])

# rag_chunkテーブルを作成

spark.sql("""

CREATE TABLE IF NOT EXISTS default.rag_chunk (

id STRING NOT NULL,

content STRING NOT NULL,

title STRING,

updated_at STRING

) USING DELTA

""")

# CDFを有効にしておきます

spark.sql("""

ALTER TABLE default.rag_chunk SET TBLPROPERTIES (delta.enableChangeDataFeed = true)

""")

CHUNK_SIZE = 1000

text_splitter = CharacterTextSplitter(chunk_size=CHUNK_SIZE, separator="\n")

# 各テキストに対してチャンク化

chunks = []

start_date = datetime(2025, 1, 1)

date_counter = 0

for doc in docs:

text_chunks = text_splitter.split_text(doc["content"])

for chunk in text_chunks:

updated_at = (start_date + timedelta(days=date_counter)).strftime("%Y-%m-%d")

chunks.append({

"id": str(uuid.uuid4()),

"content": chunk,

"title": doc["filename"],

"updated_at": updated_at,

})

date_counter += 1

# Spark DataFrameに変換

df_chunks = spark.createDataFrame(chunks, schema=schema)

# rag_chunkテーブルに追加

df_chunks.write.format("delta").mode("append").saveAsTable("default.rag_chunk")

display(df_chunks)

eval_qa = []

for doc in docs:

prompt = f"""

Based on the following text, generate 10 pairs of questions and answers that test understanding of the content.

Please use the following output format:

Q:

A:

Text:

{doc["content"]}

"""

response = client.chat.completions.create(

# Provide a valid model name for your LLM provider.

model=model_name,

messages=[{"role": "user", "content": prompt}],

temperature=0,

)

qna_text = response.choices[0].message.content

print(qna_text)

# Q&A をRAG用評価データセット形式に整形

current_q = None

current_a = None

for block in qna_text.split("\n"):

if block.startswith("Q:"):

current_q = block.replace("Q:", "").strip()

elif block.startswith("A:"):

current_a = block.replace("A:", "").strip()

if current_q and current_a:

eval_qa.append({

"inputs": {"query": current_q},

"expectations": {"expected_facts": [current_a]},

})

current_q = None

current_a = None

print(eval_qa)

import json

with open("/Volumes/workspace/default/rag_eval_qa/rag_eval_qa.json", "w", encoding="utf-8") as f:

json.dump(eval_qa, f, ensure_ascii=False, indent=2)

print("RAG用評価データセットを作成しました。")

import mlflow

from databricks.sdk import WorkspaceClient

mlflow.openai.autolog()

mlflow.set_experiment("/path/to/MLflow3")

from databricks.vector_search.client import VectorSearchClient

vs_client = VectorSearchClient()

index = vs_client.get_index(endpoint_name="vs_endpoint", index_name="workspace.default.vs_index")

from mlflow.entities import Document

from typing import List

@mlflow.trace(span_type="RETRIEVER")

def retrieve_docs(query: str) -> List[Document]:

    results = index.similarity_search(

        query_text=query,

        columns=["content"],

        num_results=10,

    )

    docs = []

    data_array = results["result"]["data_array"]

    for i, row in enumerate(data_array):

        content = row[0]

        score = row[1]

        docs.append(

            Document(

                id=f"vs_doc_{i+1}",

                page_content=content,

                metadata={"score": score},

            )

        )

    return docs

@mlflow.trace

def rag_app(query: str):

    retrieved_documents = retrieve_docs(query=query)

    retrieved_docs_text = "\n".join([f"{i+1}. {doc.page_content}" for i, doc in enumerate(retrieved_documents)])

    messages_for_llm = [

        {

            "role": "system",

            "content": f"You are an excellent assistant at 〇△× Corporation. Based on the following context, please answer questions regarding the company: {retrieved_docs_text}. If the information necessary to answer is not included in the context and the question is not of a general nature, please respond with: 'No information was found to answer the question'.",

        },

        {

            "role": "user",

            "content": query,

        },

    ]

    response = client.chat.completions.create(

        model=model_name,

        messages=messages_for_llm,

        temperature=0,

    )

    return response

import json

with open("/Volumes/workspace/default/rag_eval_qa/rag_eval_qa.json", "r") as f:

    eval_qa = json.load(f)

from mlflow.genai.scorers import (

    Correctness,

    RelevanceToQuery,

    RetrievalRelevance,

    RetrievalSufficiency,

)

# Run predefined scorers that require ground truth

mlflow.genai.evaluate(

    data=eval_qa,

    predict_fn=rag_app,

    scorers=[

        Correctness(),

        RelevanceToQuery(),

        RetrievalRelevance(),

        RetrievalSufficiency(),

    ],

)

from databricks.vector_search.reranker import DatabricksReranker

results = index.similarity_search(

        query_text=query,

        columns=["id", "content", "title", "updated_at"],

        num_results=5,

        reranker=DatabricksReranker(columns_to_rerank=["content"])

)

from databricks.vector_search.reranker import DatabricksReranker

results = index.similarity_search(

        query_text=query,

        columns=["id", "content", "title", "updated_at"],

        num_results=5,

        reranker=DatabricksReranker(columns_to_rerank=["content", "title"])

)