到目前为止,在我们的所有方法中,我们使用了嵌入模型(+词汇搜索)来识别数据集中最相关的前k个块。块的数量(k)一直是一个较小的数字,因为我们发现添加太多的块并没有帮助,并且我们的LLM有限制的上下文长度。然而,这一切都是基于这样的假设:检索到的前k个块确实是最相关的块,并且它们的顺序也是正确的。如果增加块的数量没有帮助,是因为一些相关的块在有序列表中的位置较低呢?而且,语义表示虽然非常丰富,但没有针对这个特定任务进行训练。
在本节中,我们实现了重排名,这样我们就可以使用我们的语义和词汇搜索方法在我们的数据集上投放更广泛的网(检索许多块),然后根据用户的查询重新排列顺序。这里的直觉是,我们可以通过针对我们用例的排名来弥补我们语义表示中的缺口。我们将训练一个监督模型,预测我们的文档哪一部分对给定用户的查询最相关。然后我们将使用这个预测来重新排列相关的块,使得这部分文档的块移到列表的顶部。
数据集
我们将重用我们在微调部分创建的QA数据集,因为该数据集包含了与特定章节对应的问题。我们将创建一个名为text的特征,它将章节标题和问题连接起来。我们将使用这个特征作为输入,让我们的模型预测适当的内容。我们添加章节标题(即使这些信息在来自用户查询的推理过程中不会可用),以便我们的模型可以学习如何表示用户查询中会出现的关键令牌。
def get_tag(url):
return re.findall(r"docs\.ray\.io/en/master/([^/]+)", url)[0].split("#")[0]
# Load data
from pathlib import Path
df = pd.read_json(Path(ROOT_DIR, "datasets", "embedding_qa.json"))
df["tag"] = df.source.map(get_tag)
df["section"] = df.source.map(lambda source: source.split("/")[-1])
df["text"] = df["section"] + " " + df["question"]
df.sample(n=5)
# Map only what we want to keep
tags_to_keep = ["rllib", "tune", "train", "cluster", "ray-core", "data", "serve", "ray-observability"]
df["tag"] = df.tag.apply(lambda x: x if x in tags_to_keep else "other")
Counter(df.tag)
Counter({'rllib': 1269,
'tune': 979,
'train': 697,
'cluster': 690,
'data': 652,
'ray-core': 557,
'other': 406,
'serve': 302,
'ray-observability': 175})
预处理
我们将从创建一些预处理函数开始,以更好地表示我们的数据。例如,我们的文档中有许多变量使用了驼峰命名法(例如RayDeepSpeedStrategy)。当在此上使用分词器时,我们经常丢失我们知道有用的单独令牌,而是创建了随机的子令牌。
注意:我们并非全知全能地知道要创建这些独特的预处理函数!这一切都是方法论迭代的结果。我们训练一个模型 → 观察不正确的数据点 → 查看数据如何表示(例如子令牌化)→ 更新预处理 → 迭代 ?。
import re
from transformers import BertTokenizer
def split_camel_case_in_sentences(sentences):
def split_camel_case_word(word):
return re.sub("([a-z0-9])([A-Z])", r"\1 \2", word)
processed_sentences = []
for sentence in sentences:
processed_words = []
for word in sentence.split():
processed_words.extend(split_camel_case_word(word).split())
processed_sentences.append(" ".join(processed_words))
return processed_sentences
def preprocess(texts):
texts = [re.sub(r'(?<=\w)([?.,!])(?!\s)', r' \1', text) for text in texts]
texts = [text.replace("_", " ").replace("-", " ").replace("#", " ").replace(".html", "").replace(".", " ") for text in texts]
texts = split_camel_case_in_sentences(texts) # camelcase
texts = [tokenizer.tokenize(text) for text in texts] # subtokens
texts = [" ".join(word for word in text) for text in texts]
return texts
print (preprocess(["RayDeepSpeedStrategy"]))
print (preprocess(["What is the default batch_size for map_batches?"]))
['ray deep speed strategy']
['what is the default batch size for map batch ##es ?']
训练
现在我们将训练一个简单的逻辑回归模型,该模型将根据输入文本预测标签。
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import FunctionTransformer
# Train classifier
from rag.rerank import preprocess # for pickle
reranker = Pipeline([
("preprocess", FunctionTransformer(preprocess)),
("vectorizer", TfidfVectorizer(lowercase=True)),
("classifier", LogisticRegression(multi_class="multinomial", solver="lbfgs"))
])
reranker.fit(train_df["text"].tolist(), train_df["tag"].tolist())
注意:我们还训练了一个BERT分类器,虽然其性能优于我们的逻辑分类器,但这些大型网络容易过度自信,我们不能像下面这样使用基于阈值的方法。而且没有阈值方法(我们只在重排器真正有信心时才重排),我们应用的质量分数就不会提高。
# Inference
question = "training with deepspeed"
custom_predict([question], classifier=reranker)[0]
'train'
我们现在准备评估我们训练好的重排模型。我们将使用一个自定义的预测函数,除非最高类的概率超过一定阈值,否则将预测为“其他”。
def custom_predict(inputs, classifier, threshold=0.3, other_label="other"):
y_pred = []
for item in classifier.predict_proba(inputs):
prob = max(item)
index = item.argmax()
if prob >= threshold:
pred = classifier.classes_[index]
else:
pred = other_label
y_pred.append(pred)
return y_pred
# Evaluation
metrics = {}
y_test = test_df["tag"]
y_pred = custom_predict(inputs=test_df["question"], classifier=reranker)
{
"precision": 0.9168129573272782,
"recall": 0.9171029668411868,
"f1": 0.9154520876579969,
"num_samples": 1146.0
}
测试
除了基于指标的评估之外,我们还希望评估我们的模型在一些最基本的功能测试上的表现。无论我们使用何种类型的模型,我们都需要通过所有这些基本的健全性检查。
# Basic tests
tests = [
{"question": "How to train a train an LLM using DeepSpeed?", "tag": "train"},
...
{"question": "How do I set a maximum episode length when training with Rllib", "tag": "rllib"}]
for test in tests:
question = test["question"]
prediction = predict_proba(question=test["question"], classifier=reranker)[0][1]
print (f"[{prediction}]: {question} → {preprocess([question])}")
assert (prediction == test["tag"])
[train]: How to train a train an LLM using DeepSpeed? → ['how to train a train an ll ##m using deep speed ?']
...
[rllib]: How do I set a maximum episode length when training with Rllib → ['how do i set a maximum episode length when training with r ##lli ##b']
重排实验
现在我们准备使用以下步骤在检索后应用我们的重排模型:
- 增加检索到的上下文(可以对此进行实验),以便我们可以应用重排以产生较小的子集(num_chunks)。这里的直觉是,我们将使用语义和词汇搜索检索N个块(N > k),然后我们将使用重排来重新排序检索结果(前k个)。
- 如果预测标签高于阈值,那么我们将把该标签的所有检索源移动到顶部。如果预测标签低于阈值,则不进行重排。这里的直觉是,除非我们对特定查询涉及我们文档的哪些部分(或如果它恰好涉及多个部分)有信心,否则我们不会错误地重新排列结果。
- 使用前k个检索到的块进行生成。
我们将直接修改我们的QueryAgent类以包含重排:
class QueryAgent():
def __init__(rerank=True, **kwargs):
# Reranker
self.reranker = None
if rerank:
reranker_fp = Path(EFS_DIR, "reranker.pkl")
with open(reranker_fp, "rb") as file:
self.reranker = pickle.load(file)
def __call__(rerank_threshold=0.3, rerank_k=7, **kwargs):
# Rerank
if self.reranker:
predicted_tag = custom_predict(
inputs=[query], classifier=self.reranker, threshold=rerank_threshold)[0]
if predicted_tag != "other":
sources = [item["source"] for item in context_results]
reranked_indices = get_reranked_indices(sources, predicted_tag)
context_results = [context_results[i] for i in reranked_indices]
context_results = context_results[:rerank_k]
有了这个,让我们在评估运行中使用增加了重排功能的查询代理。我们将尝试各种重排阈值。
注意:阈值为零等同于不使用任何阈值。
# Experiment
rerank_threshold_list = [0, 0.3, 0.5, 0.7, 0.9]
use_reranking = True
for rerank_threshold in rerank_threshold_list:
experiment_name = f"rerank-{rerank_threshold}"
experiment_names.append(experiment_name)
run_experiment(
experiment_name=experiment_name,
num_chunks=30, # increased num chunks since we will retrieve top k
rerank_k=NUM_CHUNKS + LEXICAL_SEARCH_K, # subset of larger num_chunks
**kwargs)
original_num_chunks = NUM_CHUNKS
NUM_CHUNKS = 30
USE_RERANKING = True
RERANK_THRESHOLD = 0.5
RERANK_K = original_num_chunks + LEXICAL_SEARCH_K
注意:在重排方面还有很多可以尝试的(增加初始的num_chunks数量,重排后添加词汇搜索结果,加权重排,其中我们提升前N个类别等)。
作为参考,这里是迄今为止的前三个实验:
[('gpt-4-1106-preview',
{'retrieval_score': 0.7288135593220338, 'quality_score': 4.209039548022599}),
('rerank-0.5',
{'retrieval_score': 0.7062146892655368,
'quality_score': 3.9519774011299433}),
('prompt-ignore-contexts',
{'retrieval_score': 0.7344632768361582, 'quality_score': 3.943502824858757})]