Reranking dengan model embedding

Kode Python untuk reranking RAG

Konten Halaman

Reranking adalah langkah kedua dalam Retrieval Augmented Generation (RAG) systems, yang berada tepat di antara Retrieving dan Generating.

Kubus listrik di ruang digital

Gambar di atas adalah bagaimana Flux-1 dev membayangkan Kubus listrik di ruang digital.

Retrieval dengan reranking

Jika kita menyimpan dokumen dalam bentuk embedding di dalam vector DB sejak awal - retrieval akan memberi kita daftar dokumen yang mirip langsung.

Reranking mandiri

Namun, jika kita terlebih dahulu mengunduh dokumen dari internet, respons sistem pencarian bisa terpengaruh oleh preferensi/algoritma penyedia pencarian, konten yang didukung, optimisasi SEO, dll. sehingga kita perlu melakukan reranking setelah pencarian.

Apa yang saya lakukan -

  • mendapatkan embedding untuk query pencarian
  • mendapatkan embedding untuk setiap dokumen. dokumen tidak diharapkan lebih dari 8k token
  • menghitung kesamaan antara query dan embedding setiap dokumen
  • mengurutkan dokumen berdasarkan kesamaan ini.

Tidak ada vector DB di sini, mari kita lanjutkan.

Contoh kode

Menggunakan Langchain untuk terhubung ke Ollama dan fungsi cosine_similarity dari langchain. Anda dapat menyaring berdasarkan ukuran kesamaan, tetapi ingat bahwa untuk domain dan embedding LLMs yang berbeda, ambang batas akan berbeda.

Saya akan senang jika bagian kode ini berguna bagi Anda dalam cara apa pun. Lisensi Copy/Paste/UseAnyWayYouWant. Terima kasih.

from langchain_core.documents import Document
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.utils.math import cosine_similarity
import numpy as np


def cosine_distance(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    return 1.0 - cosine_similarity(a, b)

def compute_score(vectors: np.ndarray) -> float:
    score = cosine_distance(vectors[0].reshape(1, -1), vectors[1].reshape(1, -1)).item()
    return score

def list_to_array(lst):
    return np.array(lst, dtype=float)   

def compute_scorel(lists) -> float:
    v1 = list_to_array(lists[0])
    v2 = list_to_array(lists[1])
    return compute_score([v1, v2])

def filter_docs(emb_model_name, docs, query, num_docs):
    content_arr = [doc.page_content for doc in docs]

    ollama_emb = OllamaEmbeddings(
        model=emb_model_name
    )

    docs_embs = ollama_emb.embed_documents(content_arr)
    query_embs = ollama_emb.embed_query(query)
    sims = []
    for i, emb in enumerate(docs_embs):
        idx = docs[i].id
        s = compute_scorel([query_embs, docs_embs[i]])
        simstr = str(round(s, 4))
        docs[i].metadata["sim"] = simstr
        sim = {
            "idx": idx,
            "i": i,
            "sim": s,
        }
        sims.append(sim)

    sims.sort(key=sortFn)

    sorted_docs = [docs[x["i"]] for x in sims]
    filtered_docs = sorted_docs[:num_docs]
    return filtered_docs

Model Embedding Terbaik

Untuk tugas saya, model embedding terbaik saat ini adalah bge-large:335m-en-v1.5-fp16

Model kedua adalah nomic-embed-text:137m-v1.5-fp16 dan jina/jina-embeddings-v2-base-en:latest.

Namun, lakukan uji coba sendiri untuk domain dan query Anda sendiri.

Tautan yang berguna