WordPress のデータを使って簡単な RAG を実装する(2)

ollama run llama3.2

CREATE EXTENSION IF NOT EXISTS vector;

pip install -r requirements.txt

import pandas as pd

df = pd.read_json('path/to/posts.json')

from langchain_community.document_loaders import DataFrameLoader
from langchain_community.document_transformers import Html2TextTransformer

df['title_content'] = df['title']['rendered'] + ' ' + df['content']['rendered']

post_loader = DataFrameLoader(df, page_content_column='title_content')
docs = post_loader.load()

html2text = Html2TextTransformer()
text_docs = html2text.transform_documents(docs)

from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores.pgvector import PGVector

pg_connection_string = PGVector.connection_string_from_db_params(
    driver="psycopg2",
    host="localhost",
    port=5432,
    database="wordpressrag",
    user="postgres",
    password="postgres",
)

embeddings = HuggingFaceEmbeddings(model_name="paraphrase-xlm-r-multilingual-v1")

db = PGVector.from_documents(
    embedding=embeddings,
    documents = text_docs,
    collection_name='example_com', # 何でも良いです
    connection_string=pg_connection_string,
)

wordpressrag=# \d
                  List of relations
 Schema |          Name           | Type  |  Owner
--------+-------------------------+-------+----------
 public | langchain_pg_collection | table | postgres
 public | langchain_pg_embedding  | table | postgres
(2 rows)

wordpressrag=# select * from langchain_pg_collection;
     name      | cmetadata |                 uuid
---------------+-----------+--------------------------------------
 example_com   | null      | 53fb1fc1-f9e6-4cbf-ada1-d77c413554e6
(1 rows)

wordpressrag=# \d langchain_pg_embedding;
               Table "public.langchain_pg_embedding"
    Column     |       Type        | Collation | Nullable | Default
---------------+-------------------+-----------+----------+---------
 collection_id | uuid              |           |          |
 embedding     | vector            |           |          |
 document      | character varying |           |          |
 cmetadata     | json              |           |          |
 custom_id     | character varying |           |          |
 uuid          | uuid              |           | not null |
Indexes:
    "langchain_pg_embedding_pkey" PRIMARY KEY, btree (uuid)
Foreign-key constraints:
    "langchain_pg_embedding_collection_id_fkey" FOREIGN KEY (collection_id) REFERENCES langchain_pg_collection(uuid) ON DELETE CASCADE

store = PGVector(
    collection_name='example_com',
    connection_string=pg_connection_string, # 前項のプログラムと同じ定義
    embedding_function=embeddings, # 前項のプログラムと同じ定義
)
retriever = store.as_retriever(search_kwargs={"k": 4}) # デフォルトが4、変更したければここを変える

from langchain_ollama import OllamaLLM

llm = OllamaLLM(model="llama3.2", num_ctx=131072)

from langchain_core.prompts import PromptTemplate

template = """以下のコンテキストのみを参照して質問に答えてください。
コンテキスト: {context}

質問: {question}
"""

custom_rag_prompt = PromptTemplate.from_template(template)

from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

def format_docs(docs):
    """文書を連結して返す。コンテキスト用の文字列を生成するために使う。"""
    return "\n\n".join(doc.page_content for doc in docs)


rag_chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()}
    | custom_rag_prompt
    | llm
    | StrOutputParser()
)

result = rag_chain.invoke("WordPress のデータを使って RAG を実装する方法を教えて。"):

# 入力は {"docs": retriever から取得した文書の配列, "question": 質問文そのまま}
prompt_input_builder = {
    "question": itemgetter("question"),
    "context": lambda x: format_docs(x["docs"]),
}

# 入力は {"docs": retriever から取得した文書の配列, "question": 質問文そのまま}
answer_chain = {
    "answer": prompt_input_builder | custom_rag_prompt | model | StrOutputParser(),
    "docs": lambda x: x["docs"]
}

rag_chain = (
    {"docs": retriever, "question": RunnablePassthrough()}
    | answer_chain
)

result = rag_chain.invoke("WordPress のデータを使って RAG を実装する方法を教えて。"):
# result['answer'] には回答が入る
# result['docs'] にはベクトルデータベースからの結果が入る

# ollama show llama3.2
  Model
    architecture        llama
    parameters          3.2B
    context length      131072
    embedding length    3072
    quantization        Q4_K_M

  Parameters
    stop    "<|start_header_id|>"
    stop    "<|end_header_id|>"
    stop    "<|eot_id|>"

  License
    LLAMA 3.2 COMMUNITY LICENSE AGREEMENT
    Llama 3.2 Version Release Date: September 25, 2024