プロンプトエンジニアリング手法 外部データ接続・RAG編

docker run -p 6379:6379 redislabs/redisearch:latest

import pandas as pd
from pyarrow.parquet import ParquetDataset

ds = ParquetDataset(
    "./shopping_queries_dataset_products.parquet",
    filters=[('product_locale', '=', 'jp')]  # 日本の商品に絞っています
)
df = ds.read().to_pandas()

import os

import tiktoken
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter

OPENAI_API_KEY = "sk-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX"

os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY

tiktoken.encoding_for_model(''text-embedding-ada-002'')
tokenizer = tiktoken.get_encoding('cl100k_base')


def tiktoken_len(text):
    tokens = tokenizer.encode(
        text,
        disallowed_special=(), )
    return len(tokens)


text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=400,
    chunk_overlap=20,
    length_function=tiktoken_len,
    separators=["nn", "n", " ", ""], )

embedding = OpenAIEmbeddings(
    model="text-embedding-ada-002",
    openai_api_key=OPENAI_API_KEY, )

import redis
from redis.commands.search.field import NumericField, TextField, VectorField
from redis.commands.search.indexDefinition import IndexDefinition, IndexType


INDEX_NAME = "rag-index"
DOC_PREFIX = "doc:"
VECTOR_DIMENSION = 1536  # text-embedding-ada-002の次元数


r = redis.Redis(host="localhost", port=6379)

def create_index(vector_dimension: int):
    try:
        r.ft(INDEX_NAME).info()
        print("Index already exists!")
    except:
        schema = (
            VectorField(
                "Embedding",
                "FLAT",
                {
                    "TYPE": "FLOAT32",
                    "DIM": vector_dimension,
                    "DISTANCE_METRIC": "COSINE",
                },
            ),
            TextField("ProductId"),
            TextField("ProductTitle"),
            TextField("ProductDescription"),
            TextField("ProductBulletPoint"),
            NumericField("ChunkNumber"),
            NumericField("ChunkCount"),
        )
        definition = IndexDefinition(
            prefix=[DOC_PREFIX],
            score_field="Embedding",
            index_type=IndexType.HASH,
        )
        r.ft(INDEX_NAME).create_index(
            fields=schema,
            definition=definition,
        )


create_index(vector_dimension=VECTOR_DIMENSION)

import json
from uuid import uuid4

import numpy as np


pipe = r.pipeline()

for i, record in df.iterrows():
    texts = text_splitter.split_text(
        (record["product_title"] or "") + (record["product_description"] or ""))
    vectors = embedding.embed_documents(texts)
    for j, (vector, text) in enumerate(zip(vectors, texts), start=1):
        mapping = {
            "ProductId": record["product_id"],
            "ProductTitle": record["product_title"],
            "ProductDescription": record["product_description"] or "",
            "ProductBulletPoint": record["product_bullet_point"] or "",
            "Embedding": np.array(vector).astype(np.float32).tobytes(),
            "ChunkNumber": j,
            "ChunkCount": len(texts),
            "content": record["product_title"],
            "metadata": json.dumps({
                "product_brand": record["product_brand"],
                "product_color": record["product_color"] or ""})}
        pipe.hset(f"doc:{uuid4()}", mapping=mapping)
pipe.execute()

from langchain.vectorstores.redis import Redis as RedisVectorstore
from langchain.chains import RetrievalQA
from langchain.chat_models import ChatOpenAI


vector_store = RedisVectorstore.from_existing_index(
    embedding=embedding,
    index_name=INDEX_NAME,
    vector_key="Embedding",
    redis_url="redis://localhost:6379", )
retriever = vector_store.as_retriever()


chain = RetrievalQA.from_chain_type(
    ChatOpenAI(),
    chain_type="stuff",
    retriever=vector_store.as_retriever(), )

> df.sample().product_title.values[0]
'What Say カラーレンズ ミラーレンズ サングラス クリアレンズ 伊達メガネ 全17色 クラシックフレーム トレンド UV400 メンズ レディース ソフト ＆ ハードケース 付 (ライト ブルー/ブラック)'

> chain.run("カラーレンズについて教えてください。")
' カラーレンズは17色から選択でき、クラシックフレームの伊達メガネを装着できます。UV400保護のレンズになっており、ソフト＆ハードケースが付属しています。'