Membangun RAG dengan Milvus dan Unstructured

Unstructured menyediakan platform dan alat untuk menelan dan memproses dokumen tak terstruktur untuk Retrieval Augmented Generation (RAG) dan penyempurnaan model. Platform ini menawarkan platform UI tanpa kode dan layanan API tanpa server, yang memungkinkan pengguna untuk memproses data pada sumber daya komputasi yang dihosting oleh Unstructured.

Dalam tutorial ini, kita akan menggunakan Unstructured untuk mengambil dokumen PDF dan kemudian menggunakan Milvus untuk membangun pipeline RAG.

Persiapan

Ketergantungan dan Lingkungan

$ pip install -qU "unstructured[pdf]" pymilvus milvus-lite openai

import os

os.environ["OPENAI_API_KEY"] = "sk-***********"

Menyiapkan klien Milvus dan OpenAI

Anda dapat menggunakan klien Milvus untuk membuat koleksi Milvus dan memasukkan data ke dalamnya.

from pymilvus import MilvusClient, DataType

# Initialize Milvus client
milvus_client = MilvusClient(uri="./milvus_demo.db")

Periksa apakah koleksi sudah ada dan hapus jika sudah ada.

collection_name = "my_rag_collection"

if milvus_client.has_collection(collection_name):
    milvus_client.drop_collection(collection_name)

Siapkan klien OpenAI untuk menghasilkan embedding dan menghasilkan respons.

from openai import OpenAI

openai_client = OpenAI()


def emb_text(text):
    return (
        openai_client.embeddings.create(input=text, model="text-embedding-3-small")
        .data[0]
        .embedding
    )

Hasilkan embedding uji coba dan cetak dimensi dan beberapa elemen pertama.

test_embedding = emb_text("This is a test")
embedding_dim = len(test_embedding)
print(embedding_dim)
print(test_embedding[:10])

1536
[0.009889289736747742, -0.005578675772994757, 0.00683477520942688, -0.03805781528353691, -0.01824733428657055, -0.04121600463986397, -0.007636285852640867, 0.03225184231996536, 0.018949154764413834, 9.352207416668534e-05]

Membuat Koleksi Milvus

Kita akan membuat sebuah koleksi dengan skema berikut ini:

• id: kunci utama, yang merupakan pengenal unik untuk setiap dokumen.
• vector: penyematan dokumen.
• text: konten teks dari dokumen.
• metadata: metadata dari dokumen.

Kemudian kita membangun sebuah indeks AUTOINDEX pada bidang vector. Dan kemudian buatlah koleksinya.

# Create schema
schema = milvus_client.create_schema(auto_id=False, enable_dynamic_field=False)
# Add fields to schema
schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True)
schema.add_field(field_name="vector", datatype=DataType.FLOAT_VECTOR, dim=embedding_dim)
schema.add_field(field_name="text", datatype=DataType.VARCHAR, max_length=65535)
schema.add_field(field_name="metadata", datatype=DataType.JSON)
index_params = MilvusClient.prepare_index_params()
index_params.add_index(
    field_name="vector",
    metric_type="COSINE",
    index_type="AUTOINDEX",
)
milvus_client.create_collection(
    collection_name=collection_name,
    schema=schema,
    index_params=index_params,
    consistency_level="Bounded",
)

milvus_client.load_collection(collection_name=collection_name)

Memuat data dari Unstructured

Unstructured menyediakan pipeline ingestion yang fleksibel dan kuat untuk memproses berbagai jenis file, termasuk PDF, HTML, dan banyak lagi. Kita akan mempartisi dan memotong file PDF lokal. Dan kemudian memuat data tersebut ke dalam Milvus.

import warnings
from unstructured.partition.auto import partition

warnings.filterwarnings("ignore")

elements = partition(
    filename="./pdf_files/WhatisMilvus.pdf",
    strategy="hi_res",
    chunking_strategy="by_title",
)  # Replace with the path to your PDF file

Mari kita periksa elemen-elemen yang dipartisi dari file PDF. Setiap elemen mewakili potongan konten yang diekstrak oleh proses partisi Unstructured.

for element in elements:
    print(element)
    break

What is Milvus?

Milvus is a high-performance, highly scalable vector database that runs efficiently across a wide range of environments, from a laptop to large-scale distributed systems. It is available as both open-source software and a cloud service.

Masukkan data ke dalam Milvus.

data = []
for i, element in enumerate(elements):
    data.append(
        {
            "id": i,
            "vector": emb_text(element.text),
            "text": element.text,
            "metadata": element.metadata.to_dict(),
        }
    )
milvus_client.insert(collection_name=collection_name, data=data)

{'insert_count': 29, 'ids': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28], 'cost': 0}

Mengambil dan Menghasilkan Respons

Tentukan fungsi untuk mengambil dokumen yang relevan dari Milvus.

def retrieve_documents(question, top_k=3):
    search_res = milvus_client.search(
        collection_name=collection_name,
        data=[emb_text(question)],
        limit=top_k,
        # search_params={"metric_type": "IP", "params": {}},
        output_fields=["text"],
    )
    return [(res["entity"]["text"], res["distance"]) for res in search_res[0]]

Tentukan fungsi untuk menghasilkan respons menggunakan dokumen yang diambil dalam pipeline RAG.

def generate_rag_response(question):
    retrieved_docs = retrieve_documents(question)
    context = "\n".join([f"Text: {doc[0]}\n" for doc in retrieved_docs])
    system_prompt = (
        "You are an AI assistant. Provide answers based on the given context."
    )
    user_prompt = f"""
    Use the following pieces of information to answer the question. If the information is not in the context, say you don't know.
    
    Context:
    {context}
    
    Question: {question}
    """
    response = openai_client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
    )
    return response.choices[0].message.content

Mari kita uji pipeline RAG dengan sebuah contoh pertanyaan.

question = "What is the Advanced Search Algorithms in Milvus?"
answer = generate_rag_response(question)
print(f"Question: {question}")
print(f"Answer: {answer}")

Question: What is the Advanced Search Algorithms in Milvus?
Answer: The Advanced Search Algorithms in Milvus include a wide range of in-memory and on-disk indexing/search algorithms such as IVF, HNSW, and DiskANN. These algorithms have been deeply optimized, and Milvus delivers 30%-70% better performance compared to popular implementations like FAISS and HNSWLib.