Эта статья представляет собой всестороннее исследование извлечения и ранжирования документов с использованием мощной структуры Llama Index. Используя возможности встраивания OpenAI и SBBERT, процесс поиска достигает высокой точности и эффективности. Кроме того, включение LLM для повторного ранжирования документов обеспечивает точную настройку результатов, оптимизируя общую производительность.

Убедитесь, что установлены все необходимые компоненты

! pip install -qqq llama_index
! pip install -qqq langchain
! pip install -qqq  sentence_transformers

Импорт всех необходимых пакетов

import nest_asyncio
nest_asyncio.apply()
import logging
import sys
import openai
from llama_index import SimpleDirectoryReader, Document
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))
from llama_index import (VectorStoreIndex,SimpleDirectoryReader,ServiceContext,LLMPredictor)
from llama_index.indices.postprocessor import LLMRerank,SentenceTransformerRerank
from llama_index.llms import OpenAI
from IPython.display import Markdown, display ,HTML
import os
import pandas as pd
pd.set_option("display.max_colwidth", -1)
from llama_index.retrievers import VectorIndexRetriever
from llama_index.indices.query.schema import QueryBundle
from copy import deepcopy
from llama_index import SimpleDirectoryReader, Document
from langchain.embeddings.huggingface import HuggingFaceEmbeddings 
from llama_index import ServiceContext, set_global_service_context,LangchainEmbedding

Для этой задачи я выбрал набор данных с открытым исходным кодом «10dataset-text-document-classification» от Kaggle, который охватывает 10 тем, связанных с бизнесом, спортом, технологиями, космосом, политикой, медициной, едой, и более.

Предварительная обработка данных

folder_path = "/kaggle/input/10dataset-text-document-classification"
folders = os.listdir(folder_path)
rows = []
for c in folders:
    docs = os.listdir(f"{folder_path}/{c}")
    for d in docs:
        with open(f"{folder_path}/{c}/{d}", "r") as file:
            rows.append([file.read(), c])
df = pd.DataFrame(rows, columns=["doc_text", "class"])
df.drop_duplicates(subset=['doc_text'],inplace=True)
df.reset_index(inplace=True,drop=True)

df['id'] = df.groupby(['doc_text']).ngroup().astype(str)
df['ID_CLASS'] = df['id']  +'_' + df['class']
df = df[['doc_text','ID_CLASS']]
df.head()

data_dict = df.set_index('ID_CLASS').to_dict()['doc_text']
documents = [Document(text=value, metadata={"filename": key}) for key, value in data_dict.items()]
documents[509]

Ниже приведен образец документа, демонстрирующий связанные метаданные
metadata={‘filename’: ‘419_medical’}

''''Document(id_='aa757b75-b2cc-40b8-ba19-8a99f6021556', 
embedding=None, 
metadata={'filename': '419_medical'}, 
excluded_embed_metadata_keys=[], 
excluded_llm_metadata_keys=[], 
relationships={}, 
hash='521ed6eab5e61bba2204dac0a83414d99aa81365a82b494225995d60de8f0929', 
text='I am writing this to find out the following:\n1.)\t
Any information on surgery to prevent reflux esophagitis.\n2.)\
tThe name(s) of a doctor(s) who specialize in such surgery.\n3.)\t
Information on reflux esophagitis which leads to cancer.\n
My boyfriend, age 34 and otherwise in good health, 
was diagnosed with \nreflux esophagitis and a hiatal h
ernia about 2 years ago.  At that time he \n
saw a gastroenterologist and has tried acid controllers (Mylanta, \nTagamet), as well as a restricted diet and raising the head of his bed.  \nThese treatments were not effective and because the damage was \nworsening, he opted for a surgical repair 3 months ago.  He was told \nthere were two repair techniques that could fix the problem; a Nissen \nwrap and a "Hill Repair".  He opted for the "Hill Repair". He recovered \nvery well from the surgery itself but the pain he had originally is worse \nand in addition he now has trouble swallowing (including saliva).\nThe doctor now wants to do an endoscopy and has also informed him \nthat a biopsy might be necessary if he has a pre-cancerous condition \nwhich he called "Barrett\'s Syndrome". If he can\'t avoid having reflux will \nhe necessarily get cancer?\nBasically, if anyone has any information on what he should do now, I\'d \nappreciate it.\nThanks,\nPat Lydon/ NetManage, Inc./ [email protected]\n', 
start_char_idx=None, 
end_char_idx=None, 
text_template='{metadata_str}\n\n{content}', 
metadata_template='{key}: {value}', 
metadata_seperator='\n')

Для встраивания наших документов мы будем использовать модель с открытым исходным кодом «sentence-transformers/all-mpnet-base-v2». Эта мощная модель позволяет нам преобразовывать текстовое содержимое наших документов в плотные многомерные векторы, фиксируя их семантическое значение.

embed_model = LangchainEmbedding(HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2"))

# Creating Chunks of documents
service_context = ServiceContext.from_defaults(
    embed_model=embed_model,
    chunk_size=128,
    chunk_overlap=10,
)

# Buid Vector Index over document chunks 
index = VectorStoreIndex.from_documents(documents, service_context=service_context,show_progress = True)

Теперь давайте приступим к этапу поиска и повторного ранжирования. Первоначально мы будем использовать реализацию с открытым исходным кодом для этой задачи. Позже мы рассмотрим возможность выполнения этого с помощью OpenAI.

def get_retrieved_nodes(query, vector_top_k=int, reranker_top_n=None, with_reranker_sbert=False,with_reranker_open_ai = False):
    query_bundle = QueryBundle(query)
    # configure retriever
    retriever = VectorIndexRetriever(index=index, similarity_top_k=vector_top_k)
    retrieved_nodes = retriever.retrieve(query_bundle)
    # configure reranker and rerank retrived nodes
    if with_reranker_sbert:
  
        reranker = SentenceTransformerRerank(model="cross-encoder/ms-marco-TinyBERT-L-2-v2",top_n=reranker_top_n)#cross-encoder/ms-marco-TinyBERT-L-2-v2 #cross-encoder/ms-marco-MiniLM-L-2-v2
        retrieved_nodes = reranker.postprocess_nodes(retrieved_nodes,query_bundle)
    
    #open ai 
    if with_reranker_open_ai:
        reranker = LLMRerank(choice_batch_size=5, top_n=reranker_top_n, service_context=service_context)
        retrieved_nodes = reranker.postprocess_nodes(retrieved_nodes, query_bundle)

    return retrieved_nodes
def pretty_print(df):
    return display(HTML(df.to_html().replace("\\n", "<br>")))
def visualize_retrieved_nodes(nodes) -> None:
    result_dicts = []
    for node in nodes:
        node = deepcopy(node)
        node_page = node.node.metadata['filename']
        node_text = node.node.get_text()
        node_text = node_text.replace("\n", " ")

        result_dict = {"Score": node.score, "Text": node_text,"page_no":node_page}
        result_dicts.append(result_dict)

    pretty_print(pd.DataFrame(result_dicts))

Давайте начнем с запуска процесса поиска без выполнения какой-либо переоценки.

new_nodes = get_retrieved_nodes("Uk economy facing risks", vector_top_k=30, with_reranker_sbert=False,with_reranker_open_ai = False,reranker_top_n=None)
visualize_retrieved_nodes(new_nodes)

Теперь давайте переоценим полученные результаты с помощью Sbert Чем выше оценка, тем лучше результаты

new_nodes = get_retrieved_nodes("Uk economy facing risks", vector_top_k=30, with_reranker_sbert=True,with_reranker_open_ai = False,reranker_top_n=5)
visualize_retrieved_nodes(new_nodes)

Теперь давайте переоценим полученные результаты с помощью OpenAI. Мы будем использовать возможности OpenAI для выполнения процесса переоценки.

Процесс поиска и ранжирования LLM включает использование большой языковой модели (LLM) для определения релевантности документов или фрагментов текста заданному запросу. Приглашение ввода включает в себя набор документов-кандидатов, и цель LLM состоит в том, чтобы выбрать соответствующие документы, предоставляя оценку их релевантности на основе внутренней метрики. Используя LLM, узлы переупорядочиваются, чтобы возвращать первые N узлов в рейтинге, гарантируя, что наиболее важная информация будет иметь приоритет.

os.environ['OPENAI_API_KEY'] = "Enter-your-open-ai-key"
openai.api_key = os.environ.get('OPENAI_API_KEY')

new_nodes = get_retrieved_nodes("Uk economy facing risks", vector_top_k=30, with_reranker_sbert=False,with_reranker_open_ai = True,reranker_top_n=5)
visualize_retrieved_nodes(new_nodes)

Дополнительные действия

Удалите узлы, оценка сходства которых ниже порогового значения.

# SimilarityPostprocessor
# Used to remove nodes that are below a similarity score threshold.

from llama_index.indices.postprocessor import SimilarityPostprocessor

postprocessor = SimilarityPostprocessor(similarity_cutoff=0.7)

postprocessor.postprocess_nodes(nodes)

Включить, чтобы исключить ключевые слова для уточнения результатов

# KeywordNodePostprocessor
# Used to ensure certain keywords are either excluded or included.

from llama_index.indices.postprocessor import KeywordNodePostprocessor

postprocessor = KeywordNodePostprocessor(
  required_keywords=["word1", "word2"],
  exclude_keywords=["word3", "word4"]
)

postprocessor.postprocess_nodes(nodes)

Документация
https://gpt-index.readthedocs.io/en/latest/core_modules/supporting_modules/service_context.html