Optimizing RAG — Supervised Embeddings & Reranking โดยใช้ LlamaIndex ด้วยข้อมูลของเราเอง

Manusaporn Treerungroj

Published in

Noob Learning

10 min readMay 13, 2024

Blog นี้ถูกแปลมาจากต้นฉบับโดยคุณ Deltaaruna ถ้าพร้อมแล้ว ไปอ่านกันเล้ย

1. Introduction

เราสามารถสร้าง RAG system โดยใช้ LlamaIndex ซึ่งรองรับการทำงานของ embedding และ reranking หลายแบบ ใน blog นี้จะอธิบายว่า embedding และ reranking ทำงานยังไง และเราจะ train ตัว embedding ด้วยข้อมูลของเราเองได้ยังไง ใครที่ยังไม่รู้จัก LlamaIndex ย้อนไปอ่าน blog นี้ก่อนได้จ้า

2. Vector embeddings

LlamaIndex จะสร้าง vector embeddings จาก document มาดูกันก่อนว่าจะ optimize ตรงนี้ได้ยังไง

2.1 Bi-encoders

นึกภาพง่ายๆ ว่า bi-encoders จะสร้าง vector embeddings จาก string ที่ใส่ไป

ทีนี้มาดูกันว่า vector DB search ทำงานยังไง พอเราถามคำถามไป คำถามนั้นจะถูกแปลงเป็น vector embedding แล้วจะทำ similarity search ใน vector store เพื่อหา document ที่เกี่ยวข้อง/ใกล้เคียงมากที่สุด โดยใช้วิธีประเมินความคล้ายคลึงอย่าง cosine similarity index ซึ่งมีขั้นตอนดังนี้

สร้าง vector embeddings จากข้อมูลของเรา
แปลงคำถามให้เป็น vector embeddings
Retrieve documents ที่มี similarity index สูงสูด (แปลว่าใกล้เคียงมากสุด)

เพื่อให้เห็นภาพมากขึ้น ลองเขียน code กัน! (รันบน Colab GPU นะ)

!pip install -U sentence-transformers rank_bm25

import json
from sentence_transformers import SentenceTransformer, CrossEncoder, util
import gzip
import os
import torch
if not torch.cuda.is_available():
print("Warning: No GPU found. Please add GPU to your notebook")
#We use the Bi-Encoder to encode all passages, so that we can use it with semantic search
bi_encoder = SentenceTransformer('multi-qa-MiniLM-L6-cos-v1')
bi_encoder.max_seq_length = 256 #Truncate long passages to 256 tokens
top_k = 32 #Number of passages we want to retrieve with the bi-encoder
# As dataset, we use Simple English Wikipedia. Compared to the full English wikipedia, it has only
# about 170k articles. We split these articles into paragraphs and encode them with the bi-encoder
wikipedia_filepath = 'simplewiki-2020–11–01.jsonl.gz'
if not os.path.exists(wikipedia_filepath):
util.http_get('http://sbert.net/datasets/simplewiki-2020-11-01.jsonl.gz', wikipedia_filepath)
passages = []
with gzip.open(wikipedia_filepath, 'rt', encoding='utf8') as fIn:
for line in fIn:
data = json.loads(line.strip())
#Add all paragraphs
#passages.extend(data['paragraphs'])
#Only add the first paragraph
passages.append(data['paragraphs'][0])
print("Passages:", len(passages))
# We encode all passages into our vector space. This takes about 5 minutes (depends on your GPU speed)
corpus_embeddings = bi_encoder.encode(passages, convert_to_tensor=True, show_progress_bar=True)
Here we encode all our data into a vector space. Then lets do a semantic search in the vector embeddings
# This function will search all wikipedia articles for passages that
# answer the query
def search(query):
print("Input question:", query)
##### Semantic Search #####
# Encode the query using the bi-encoder and find potentially relevant passages
question_embedding = bi_encoder.encode(query, convert_to_tensor=True)
question_embedding = question_embedding.cuda()
hits = util.semantic_search(question_embedding, corpus_embeddings, top_k=top_k)
hits = hits[0] # Get the hits for the first query

ได้ผลลัพธ์ 3 อันดับแรกคือ…

เห็นได้ชัดเลยว่า ถึงแม้ว่า bi-encoders จะเร็ว แต่ไม่ได้แม่นยำทุกครั้ง แล้วเราจะทำยังไงให้มันดีขึ้นได้บ้าง?

2.2 Cross-encoders

มี encoder อีกประเภทนึงที่เรียกว่า cross-encoder ซึ่งทำงานต่างจาก bi-encoder ที่รับแค่ประโยคเดียวเป็น input แต่ cross-encoder จะรับสองประโยคเป็น input แล้วให้ output เป็นค่าระหว่าง 0 ถึง 1 ลองดูภาพประกอบนี้

cross-encoder สามารถคำนวณผลลัพธ์ได้แม่นยำกว่า bi-encoder เยอะเลย! งั้นก็รอดละมะ? ตอนทำ information retrieval เราก็ใช้ cross-encoder เปรียบเทียบคำถามกับประโยคทั้งหมดใน vector space แล้วหาประโยคที่ความคล้ายคลึงสูงที่สุด แน่นอนว่าวิธีนี้ให้ผลลัพธ์ที่แม่นยำ แต่ข้อเสียก็คือใช้พลังคำนวณเยอะมาก อ้าววว

2.3 Cross-encoders VS bi-encoders

สรุปเราจะทำไงดี? bi-encoder เร็ว ประหยัดพลังงาน แต่แม่นยำน้อยหน่อย ส่วน cross-encoder แม่นยำกว่าเยอะ แต่กินพลังงาน แล้วจะทำยังไงให้ได้ the better of both worlds? 🤔

การคำนวณของ cross-encoder ขึ้นอยู่กับจำนวนครั้งการเปรียบเทียบที่มันทำ ยิ่งเปรียบเทียบเยอะ ก็ยิ่งกินพลังงาน เราเลยต้องหาทางลดการเปรียบเทียบที่ไม่จำเป็นลง นั่นคือการเอาข้อมูลที่ไม่จำเป็นออกก่อน จะได้ไม่ต้องเปรียบเทียบทุกอัน

วิธีก็คือ เราจะใช้ bi-encoder ไป retrieve มาให้เบื้องต้น (ถึงจะแม่นยำน้อยกว่า) เพื่อให้ได้ข้อมูลเซ็ตเล็กๆ ที่ดีพอมา แล้วค่อยเอาไปรัน cross-encoder ต่อ แบบนี้ก็จะประหยัดพลังงาน ผลลัพธ์ก็แม่นยำเว่อร์!

# This function will search all wikipedia articles for passages that
# answer the query
def search_with_cross_encoder(query):
print("Input question:", query)
##### Semantic Search #####
# Encode the query using the bi-encoder and find potentially relevant passages
question_embedding = bi_encoder.encode(query, convert_to_tensor=True)
question_embedding = question_embedding.cuda()
hits = util.semantic_search(question_embedding, corpus_embeddings, top_k=top_k)
hits = hits[0] # Get the hits for the first query
##### Re-Ranking #####
# Now, score all retrieved passages with the cross_encoder
cross_inp = [[query, passages[hit['corpus_id']]] for hit in hits]
cross_scores = cross_encoder.predict(cross_inp)
# Sort results by the cross-encoder scores
for idx in range(len(cross_scores)):
hits[idx]['cross-score'] = cross_scores[idx]
# Output of top-5 hits from bi-encoder
print("\n - - - - - - - - - - - - -\n")
print("Top-3 Bi-Encoder Retrieval hits")
hits = sorted(hits, key=lambda x: x['score'], reverse=True)
for hit in hits[0:3]:
print("\t{:.3f}\t{}".format(hit['score'], passages[hit['corpus_id']].replace("\n", " ")))
# Output of top-5 hits from re-ranker
print("\n - - - - - - - - - - - - -\n")
print("Top-3 Cross-Encoder Re-ranker hits")
hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True)
for hit in hits[0:3]:
print("\t{:.3f}\t{}".format(hit['cross-score'], passages[hit['corpus_id']].replace("\n", " ")))

ได้ผลลัพธ์งี้

ดูดีขึ้นใช่มะ คราวนี้เราก็รู้วิธี optimize ตัว RAG retrievals แล้ว ก็คือใช้ cross-encoder มาทำ reranking อีกทีได้เล้ย

3. Augmenting encoders ด้วยข้อมูลเราเอง

ถ้าเราใช้ embedding model ที่มีประสิทธิภาพ ก็จะแม่นยำขึ้นได้ และถ้าเราสามารถ train embedding model โดยใช้ข้อมูลของเราเอง ก็จะยิ่งมีแม่นยำมากขึ้นไปอีก นี่เป็น idea ของงานวิจัย Augmented SBERT: Data Augmentation Method for Improving Bi-Encoders for Pairwise Sentence Scoring Tasks

เราสามารถใช้เทคนิคนี้เพื่อ fine-tune ตัว embedding ให้ละเอียดยิ่งขึ้นได้ ทำให้ผลลัพธ์สุดท้ายมีความแม่นยำมากขึ้น pipeline การทำงานเป็นแบบนี้

Train ตัว cross-encoder ด้วย dataset ของเราเอง (เรียกว่า “Golden-Dataset” หรือ Dataset A ในภาพ): ตัวอย่างเช่น dataset เกี่ยวกับบริษัทของเราเอง
สร้าง narrowed dataset (เรียกว่า “Silver-Dataset” หรือ Dataset B ในภาพ): dataset ใหม่นี้อาจจะเฉพาะเจาะจงไปที่แผนกใดแผนกหนึ่งภายในบริษัท
ใช้ cross-encoder ที่ train ไว้ในขั้นที่ 1 ในการ label “Silver-Dataset”
เอา Silver-Dataset ที่ label แล้วไป train ตัว bi-encoder ต่อ
ใช้ bi-encoder เพื่อ encode data

โอเคร พักเรื่อง LlamaIndex ไว้ก่อน ตอนนี้เรามา focus ที่การสร้าง pipeline ตามขั้นตอนข้างบนกันดีกว่า source codes ทั้งหมดมีแปะอยู่บน GitHub นะ เพื่อนๆ ลอง clone repo เล่นไปอ่านไปได้

3.1 Train ตัว cross-encoder ด้วย “Golden-Dataset” ของเราเอง

ก่อนหน้านี้เราได้พูดถึง cross-encoder กันไปแล้ว ตอนนี้เรามาดูวิธีการ train ตัว cross-encoder ด้วย dataset ที่มีอยู่แล้วกันบ้าง

เราจะ train ตัว cross-encoder แบบ supervised learning โดย dataset ที่จะใช้ประกอบไปด้วยคู่ประโยคที่มี numerical correlation คือ ค่าที่บ่งบอกถึงความคล้ายกันระหว่างสองประโยคในแต่ละคู่

from torch.utils.data import DataLoader
from sentence_transformers import models, losses, util, LoggingHandler, SentenceTransformer
from sentence_transformers.cross_encoder import CrossEncoder
from sentence_transformers.cross_encoder.evaluation import CECorrelationEvaluator
from sentence_transformers.evaluation import BinaryClassificationEvaluator
from sentence_transformers.readers import InputExample
from datetime import datetime
from zipfile import ZipFile
import logging
import csv
import sys
import torch
import math
import gzip
import os


# You can specify any huggingface/transformers pre-trained model here, for example, bert-base-uncased, roberta-base, xlm-roberta-base
model_name = "bert-base-uncased"
batch_size = 16
num_epochs = 1
max_seq_length = 128
use_cuda = torch.cuda.is_available()

###### Read Datasets ######
sts_dataset_path = "datasets/stsbenchmark.tsv.gz"
qqp_dataset_path = "quora-IR-dataset"


# Check if the STSb dataset exists. If not, download and extract it
if not os.path.exists(sts_dataset_path):
    util.http_get("https://sbert.net/datasets/stsbenchmark.tsv.gz", sts_dataset_path)

# Check if the QQP dataset exists. If not, download and extract
if not os.path.exists(qqp_dataset_path):
    logging.info("Dataset not found. Download")
    zip_save_path = "quora-IR-dataset.zip"
    util.http_get(url="https://sbert.net/datasets/quora-IR-dataset.zip", path=zip_save_path)
    with ZipFile(zip_save_path, "r") as zipIn:
        zipIn.extractall(qqp_dataset_path)

# Cross encoder will be saved in this location. Path will be like this output/cross-encoder/stsb_indomain_bert-base-uncased
# Knowing the path correctly will be helpful in next step
cross_encoder_path = (
    "output/cross-encoder/stsb_indomain_"
    + model_name.replace("/", "-")
)

# Bi encoder will be saved in this location. Path will be like this output/cross-encoder/qqp_cross_domain_bert-base-uncased
# Knowing the path correctly will be helpful in next step
bi_encoder_path = (
    "output/bi-encoder/qqp_cross_domain_"
    + model_name.replace("/", "-")
)

###### Cross-encoder (simpletransformers) ######

logging.info("Loading cross-encoder model: {}".format(model_name))
# Use Huggingface/transformers model (like BERT, RoBERTa, XLNet, XLM-R) for cross-encoder model
cross_encoder = CrossEncoder(model_name, num_labels=1)

###### Bi-encoder (sentence-transformers) ######

logging.info("Loading bi-encoder model: {}".format(model_name))

# Use Huggingface/transformers model (like BERT, RoBERTa, XLNet, XLM-R) for mapping tokens to embeddings
word_embedding_model = models.Transformer(model_name, max_seq_length=max_seq_length)

# Apply mean pooling to get one fixed sized sentence vector
pooling_model = models.Pooling(
    word_embedding_model.get_word_embedding_dimension(),
    pooling_mode_mean_tokens=True,
    pooling_mode_cls_token=False,
    pooling_mode_max_tokens=False,
)

bi_encoder = SentenceTransformer(modules=[word_embedding_model, pooling_model])


# Step 1: Train cross-encoder model with STSbenchmark

logging.info("Step 1: Train cross-encoder: {} with STSbenchmark (source dataset)".format(model_name))

gold_samples = []
dev_samples = []
test_samples = []

with gzip.open(sts_dataset_path, "rt", encoding="utf8") as fIn:
    reader = csv.DictReader(fIn, delimiter="\t", quoting=csv.QUOTE_NONE)
    for row in reader:
        score = float(row["score"]) / 5.0  # Normalize score to range 0 ... 1

        if row["split"] == "dev":
            dev_samples.append(InputExample(texts=[row["sentence1"], row["sentence2"]], label=score))
        elif row["split"] == "test":
            test_samples.append(InputExample(texts=[row["sentence1"], row["sentence2"]], label=score))
        else:
            # As we want to get symmetric scores, i.e. CrossEncoder(A,B) = CrossEncoder(B,A), we pass both combinations to the train set
            gold_samples.append(InputExample(texts=[row["sentence1"], row["sentence2"]], label=score))
            gold_samples.append(InputExample(texts=[row["sentence2"], row["sentence1"]], label=score))


# We wrap gold_samples (which is a List[InputExample]) into a pytorch DataLoader
train_dataloader = DataLoader(gold_samples, shuffle=True, batch_size=batch_size)


# We add an evaluator, which evaluates the performance during training
evaluator = CECorrelationEvaluator.from_input_examples(dev_samples, name="sts-dev")

# Configure the training
warmup_steps = math.ceil(len(train_dataloader) * num_epochs * 0.1)  # 10% of train data for warm-up
logging.info("Warmup-steps: {}".format(warmup_steps))

# Train the cross-encoder model
cross_encoder.fit(
    train_dataloader=train_dataloader,
    evaluator=evaluator,
    epochs=num_epochs,
    evaluation_steps=1000,
    warmup_steps=warmup_steps,
    output_path=cross_encoder_path,
)

อะะ มาดู code กันหน่อย

เราใช้โมเดลชื่อ “bert-base-uncased” โดยลองใช้ค่า batch_size = 16, num_epochs = 1 และ max_seq_length = 128 (ตอนใช้งานจริง ควรตั้งให้ค่ามากกว่านี้นะ)

model_name = "bert-base-uncased"
batch_size = 16
num_epochs = 1
max_seq_length = 128
use_cuda = torch.cuda.is_available()

ต่อไปก็ download ตัว STS benchmark dataset และ Quora Question Pair dataset

# Check if the STSb dataset exists. If not, download and extract it
if not os.path.exists(sts_dataset_path):
    util.http_get("https://sbert.net/datasets/stsbenchmark.tsv.gz", sts_dataset_path)

# Check if the QQP dataset exists. If not, download and extract
if not os.path.exists(qqp_dataset_path):
    logging.info("Dataset not found. Download")
    zip_save_path = "quora-IR-dataset.zip"
    util.http_get(url="https://sbert.net/datasets/quora-IR-dataset.zip", path=zip_save_path)
    with ZipFile(zip_save_path, "r") as zipIn:
        zipIn.extractall(qqp_dataset_path)

จากนั้นจะใช้ mean pooling เพื่อให้ได้ vector ที่ขนาดคงที่ ขั้นตอนนี้สำคัญมากสำหรับการ fine-tune ตัว cross-encoder ให้เหมาะกับงานเฉพาะด้าน เสร็จแล้วก็กำหนดตัว cross-encoder ให้เป็น SentenceTransformer โดยใช้ embedding model กับ pooling model ของเรา

# Use Huggingface/transformers model (like BERT, RoBERTa, XLNet, XLM-R) for mapping tokens to embeddings
word_embedding_model = models.Transformer(model_name, max_seq_length=max_seq_length

# Apply mean pooling to get one fixed sized sentence vector
pooling_model = models.Pooling(
    word_embedding_model.get_word_embedding_dimension(),
    pooling_mode_mean_tokens=True,
    pooling_mode_cls_token=False,
    pooling_mode_max_tokens=False,
)
bi_encoder = SentenceTransformer(modules=[word_embedding_model, pooling_model])

สุดท้ายก็กำหนดตัว data loader และ evaluator และแล้วก็พร้อม train ตัว cross-encoder ด้วย Golden-Dataset แล้ว!

# We wrap gold_samples (which is a List[InputExample]) into a pytorch DataLoader
train_dataloader = DataLoader(gold_samples, shuffle=True, batch_size=batch_size)

# We add an evaluator, which evaluates the performance during training
evaluator = CECorrelationEvaluator.from_input_examples(dev_samples, name="sts-dev")
# Configure the training
warmup_steps = math.ceil(len(train_dataloader) * num_epochs * 0.1)  # 10% of train data for warm-up
logging.info("Warmup-steps: {}".format(warmup_steps))
# Train the cross-encoder model
cross_encoder.fit(
    train_dataloader=train_dataloader,
    evaluator=evaluator,
    epochs=num_epochs,
    evaluation_steps=1000,
    warmup_steps=warmup_steps,
    output_path=cross_encoder_path,
)

เรียบร้อย~ ได้ cross-encoder ที่ผ่านการ fine-tune ด้วย Golden-Dataset แว้ว

3.2 ใช้ trained cross-encoder มา label ให้กับ Silver-Dataset

อย่างที่รู้กันว่า cross-encoder ใช้คำนวณ similarity ระหว่างคู่ประโยคได้ เพราะงั้นเดี๋ยวเราจะใช้ cross-encoder ที่ train ด้วย Golden-Dataset มา label ให้กับคู่ประโยคใน Silver-Dataset กัน

เป้าหมายของขั้นตอนนี้คือการสร้าง training data ที่มีคุณภาพ เพื่อให้เรามีข้อมูลเพียงพอสำหรับการ fine-tune ตัว bi-encoder โดยยิ่ง Silver-Dataset มีขนาดใหญ่กว่า Golden-Dataset เท่าไหร่ ก็ยิ่งดี เพราะ cross-encoder ที่มีความแม่นยำ ก็จะสร้าง training data ที่แม่นยำได้นั่นเอง

cross-encoder จะคำนวณคะแนนสำหรับคู่ประโยคแต่ละคู่ โดยคะแนนจะเป็นแบบ binary

คะแนนมากกว่า 0.5 จะถือว่าเป็น “one”
คะแนนน้อยกว่าหรือเท่ากับ 0.5 จะถือว่าเป็น “zero”

หมายเหตุ ขั้นตอนการ label ด้วย cross-encoder นี้อาจใช้เวลานานนิสนึง

# Step 2: Label QQP train dataset using cross-encoder (BERT) model

logging.info("Step 2: Label QQP (target dataset) with cross-encoder: {}".format(model_name))

cross_encoder = CrossEncoder(cross_encoder_path)

silver_data = []

with open(os.path.join(qqp_dataset_path, "classification/train_pairs.tsv"), encoding="utf8") as fIn:
    reader = csv.DictReader(fIn, delimiter="\t", quoting=csv.QUOTE_NONE)
    for row in reader:
        if row["is_duplicate"] == "1":
            silver_data.append([row["question1"], row["question2"]])

silver_scores = cross_encoder.predict(silver_data)

# All model predictions should be between [0,1]
assert all(0.0 <= score <= 1.0 for score in silver_scores)

binary_silver_scores = [1 if score >= 0.5 else 0 for score in silver_scores]

ตอนนี้เราก็มี quality dataset พร้อม train bi-encoder แล้ว!

3.3 Train bi-encoder ด้วย Silver-Dataset

# Step 3: Train bi-encoder (SBERT) model with QQP dataset - Augmented SBERT

logging.info("Step 3: Train bi-encoder: {} over labeled QQP (target dataset)".format(model_name))

# Convert the dataset to a DataLoader ready for training
logging.info("Loading BERT labeled QQP dataset")
qqp_train_data = list(
    InputExample(texts=[data[0], data[1]], label=score) for (data, score) in zip(silver_data, binary_silver_scores)
)


train_dataloader = DataLoader(qqp_train_data, shuffle=True, batch_size=batch_size)
train_loss = losses.MultipleNegativesRankingLoss(bi_encoder)

###### Classification ######
# Given (question1, question2), is this a duplicate or not?
# The evaluator will compute the embeddings for both questions and then compute
# a cosine similarity. If the similarity is above a threshold, we have a duplicate.
logging.info("Read QQP dev dataset")

dev_sentences1 = []
dev_sentences2 = []
dev_labels = []

with open(os.path.join(qqp_dataset_path, "classification/dev_pairs.tsv"), encoding="utf8") as fIn:
    reader = csv.DictReader(fIn, delimiter="\t", quoting=csv.QUOTE_NONE)
    for row in reader:
        dev_sentences1.append(row["question1"])
        dev_sentences2.append(row["question2"])
        dev_labels.append(int(row["is_duplicate"]))

evaluator = BinaryClassificationEvaluator(dev_sentences1, dev_sentences2, dev_labels)

# Configure the training.
warmup_steps = math.ceil(len(train_dataloader) * num_epochs * 0.1)  # 10% of train data for warm-up
logging.info("Warmup-steps: {}".format(warmup_steps))

# Train the bi-encoder model
bi_encoder.fit(
    train_objectives=[(train_dataloader, train_loss)],
    evaluator=evaluator,
    epochs=num_epochs,
    evaluation_steps=1000,
    warmup_steps=warmup_steps,
    output_path=bi_encoder_path,
)

# Evaluate Augmented SBERT performance on QQP benchmark dataset

# Loading the augmented sbert model
bi_encoder = SentenceTransformer(bi_encoder_path)

logging.info("Read QQP test dataset")
test_sentences1 = []
test_sentences2 = []
test_labels = []

with open(os.path.join(qqp_dataset_path, "classification/test_pairs.tsv"), encoding="utf8") as fIn:
    reader = csv.DictReader(fIn, delimiter="\t", quoting=csv.QUOTE_NONE)
    for row in reader:
        test_sentences1.append(row["question1"])
        test_sentences2.append(row["question2"])
        test_labels.append(int(row["is_duplicate"]))

evaluator = BinaryClassificationEvaluator(test_sentences1, test_sentences2, test_labels)
bi_encoder.evaluate(evaluator)

3.4 ใช้ trained bi-encoder เพื่อสร้าง vector embeddings

ตอนนี้เราเข้าใจหลักการเบื้องต้นของการ train ตัว bi-encoder แล้ว ถึงเวลาเอา trained bi-encoder และ trained cross-encoder ไปใช้กับ Llamaindex แล้ว

ใน Colab directory จะเห็น trained bi-encoder และ trained cross-encoder แบบในรูปข้างล่างนี้

นอกจากนี้ เราต้องเอาข้อมูลไป encode ก่อน ในตัวอย่างนี้เราจะสร้าง directory ชื่อว่า data ไว้ใน datasets และเพิ่ม text file ชื่อ paul_graham_essay.txt

จะได้แบบนี้

อย่าเพิ่งตุยนะ มาลองดู code กันต่อก่อน เราจะใช้ trained bi-encoder สำหรับทำ embedding และใช้ trained cross-encoder สำหรับทำ reranking

import os

from llama_index.core import VectorStoreIndex,SimpleDirectoryReader
from llama_index.core.postprocessor import SentenceTransformerRerank
import openai

openai.api_key = ""
# Load documents from a datasets/data
documents = SimpleDirectoryReader('datasets/data').load_data()

#Now lets use the bi encoder trained earlier
print("######## Using bi-encoder trained with our data")

index = VectorStoreIndex(
        documents, embed_model="local:output/bi-encoder/qqp_cross_domain_bert-base-uncased", show_progress=True
    )
query_engine = index.as_query_engine()

response = query_engine.query("What did the author do growing up?")
print(response)

#Now lets use the cross encoder trained earlier for re-ranking
print("######## Using cross-encoder trained with our data for reranking")

# Initialize the reranker
rerank = SentenceTransformerRerank(
    model="./output/cross-encoder/stsb_indomain_bert-base-uncased", top_n=3
)


# Build the query engine
query_engine = index.as_query_engine(similarity_top_k=10, node_postprocessors=[rerank])

response = query_engine.query("What did the author do growing up?")
print(response)

นอกจากการทำ reranking ด้วย cross-encoder แล้ว Llamaindex ยัง support วิธีการ reranking แบบอื่นๆ ด้วย เช่น LLMrank, Cohere rerank, Colbert rerank เพื่อนๆ ลองไปศึกษาดูตัวอย่าง llamaindex เพิ่มเติมได้จ้า

ติดตาม Linkedin ของคุณ Deltaaruna (ผู้เขียน blog นี้) วันนี้ลาไปก่อนม้า 🐴 555 เจอกันอีกทีจ้า

ส่วนใครสนใจอยากเริ่มเรียน llm ก็ตามไปดูย้อนหลังกันได้ที่นี่เลยจ้า 👇 https://medium.com/nooblearning/a8c07cd685ad

Live ทั้งหมด by Noob Learning

อัพเดท 31–03–2024 ❤️

medium.com