Привет, Хабр!

Это вторая часть серии статьей, посвященной использованию языковых моделей. В прошлой статье я рассмотрел основные элементы, с которыми приходится работать в библиотеке Langchain. Во второй части поговорим про векторные хранилища и разделители текста.

1. TextSplitters

Разбиение текста - важный этап предварительной обработки, который позволяет разделять большие тексты на более маленькие фрагменты, которые в свою очередь могут быть переданы модели.

Существуют разные стратегии разбиения текста, но самый простой из них заключается в разделении на основании длины фрагмента. Это гарантирует, что каждый элемент будет иметь приблизительно одинаковую длину.

В этом подходе можно выделить два вида:

Разделение на основе токенов
Разделение на основе символов

Разделение на основе символов

Для этого воспользуемся классом CharacterTextSplitter:

from langchain_text_splitters import CharacterTextSplitter

Создадим объект класса со следующими параметрами:

sepator - символ или набор символов по которым будет происходить разделение текста, по умолчанию имеет значение '\n\n'
chunk_size - длина одного фрагмента. При указании этого параметра нужно помнить, что оно задает лишь примерное значение. Итоговая длина может быть как больше, так и меньше.
chunk_overlap - при разделении текста мы можем разделить смысл единого куска текста. Поэтому мы используем chunk_overlap, чтобы захватить часть предыдущего фрагмента и сохранить смысл.
length_function - функция, с помощью которой мы вычисляем длину фрагмента .
is_separator_regex - указывает, является ли указанный separator регулярным выражением или конкретной строкой.

text_splitter = CharacterTextSplitter(
    separator="\n\n",
    chunk_size=500,
    chunk_overlap=0,
    length_function=len,
    is_separator_regex=False
)

Для примеров я буду использовать файл, который доступен по ссылке:

https://drive.google.com/file/d/1fn6aaOwcblZfLDu1cYBYimFwKXzVNUrL/view?usp=sharing

Он представляет собой несколько случайных текстов на английском языке.

Итак, у CharacterTextSplitter есть два метода, которыми мы можем пользоваться для решения текущей задачи:

split_text - разделяет текст и возвращает фрагменты в виде обычных строк.
create_documents - разделяет текст и возвращает список фрагментов в виде объектов Document, которые являются сущностью LangChain.

Загрузим документ:

with open("docs.txt") as f:
    text_union = f.read()

Разделим документ на фрагменты с помощью split_text. Распечатаем фрагменты с указанием их длины:

chunks = text_splitter.split_text(text_union)

for chunk in chunks:
    print(len(chunk), chunks)

В результате получим:

Created a chunk of size 646, which is longer than the specified 500

410 There are four types of schools in the English and Welsh education system - nursery. primary, secondary and private schools. Scotland has its own education system, which is different.\n
Children start school at the age of five, but there is some free nursery-school education before that age. The state nursery schools are not for all. They are for some families, for example for families with only one parent.

........

646 Secondary schools are usually much larger than primary schools and most children - over 80 percent - go to a comprehensive school at the age of 11. These schools are for all. Pupils do not need (to pass an exam to go to these schools.
These schools are large. They have from 1.200 - 2.500 pupils. School lasts all day in the UK, so there is only one shift. In some areas there are grammar schools. Pupils must pass special exams to go to these schools.
Some parents prefer private education. In England and Wales, private schools are called public schools. They are very expensive. Only 5 per cent of the school population goes to public schools.
---------------------------------

Я пропустил некоторые фрагменты для лучшей читаемости. Как видим, сначала мы получили предупреждение, что с текущим sepator были получены фрагменты с длиной больше 500 символов, затем были распечатаны фрагменты. Здесь стоит обратить внимание, что фрагменты могут иметь длину и менее 500 символов.

Этот же пример, но с использованием create_documents:

chunks = text_splitter.create_documents([text_union])

for chunk in chunks:
    print(len(chunk.page_content), chunks)

Основные отличия - в качестве значения передаем текст в виде списка, в качестве результата получаем объект Document.

Пример 2.

sepator может быть чем угодно, например, строкой 'are'. Установим separator='are' и chunk_overlap = 100.

text_splitter = CharacterTextSplitter(
    separator="are",
    chunk_size=500,
    chunk_overlap=100,
    length_function=len,
    is_separator_regex=False
)

В результате получим примерно следующее:

462 page_content='There are four types of schools in the English and Welsh education system - nursery. primary, secondary and private schools. Scotland has its own education system, which is different.\n
Children start school at the age of five, but there is some free nursery-school education before that age. The state nursery schools are not for all. They are for some families, for example for families with only one parent.

In most areas there are private nursery schools. P'
---------------------------------
478 page_content='In most areas there are private nursery schools. Parents who want their children to go to nursery school pay for their children under 5 years old to go to these private nursery schools.
Primary school is divided into infant school (pupils from 5 to 7 years old) and junior school (from 8 to 11 years old). In some areas there are middle schools instead of junior schools, which take pupils from 9 to 12 years old.
Primary schools have from 50-200 pupils.

Как видим, текст был разделен на фрагменты примерно по 500 символов по строке 'are'. Помимо этого, благодаря chunk_overlap, второй фрагмент начинается с предложения на котором закончился первый фрагмент.

Наглядный пример разделения на фрагменты и перекрытия фрагментов:

Демонстрация доступна по ссылке: https://chunkviz.up.railway.app/

А также, если хотите больше узнать про разделение текста не только в LangChain:

RetrievalTutorials/tutorials/LevelsOfTextSplitting/5_Levels_Of_Text_Splitting.ipynb at main · FullStackRetrieval-com/RetrievalTutorials

Разделение на основе токенов

В примерах выше разделение было основа на количестве символов каждого фрагмента, теперь попробуем использовать токены вместо символов. Про токенизацию можете почитать в моей статье про дообучение модели GPT2.

Сам токенизатор я буду брать с HuggingFace:

from transformers import AutoTokenizer

model_id = "microsoft/Phi-3-mini-4k-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)

Воспользуемся метод from_huggingface_tokenizer, который предоставляет CharacterTextSplitter и установим небольшое значение chunk_size:

text_splitter = CharacterTextSplitter.from_huggingface_tokenizer(
    tokenizer, chunk_size=100, chunk_overlap=0
)

Разделим текст на фрагменты, посмотрим на их количество и содержимое:

chunks = text_splitter.split_text(state_of_the_union)
print('Всего фрагментов', len(chunks))

for chunk in chunks:
    print(len(chunk), chunk)

В результате получим 5 текстов с длиной фрагментов 410-1178:

Всего фрагментов 5

410 There are four types of schools in the English and Welsh education system - nursery. primary, secondary and private schools. Scotland has its own education system, which is different.\n
Children start school at the age of five, but there is some free nursery-school education before that age. The state nursery schools are not for all. They are for some families, for example for families with only one parent.
------------------------------

Всего получили 5 фрагментов длина которых намного больше указанной.

RecursiveCharacterTextSplitter

Во всем примерах выше были фрагменты текста, которые превышали максимальную длину. Но у нас может возникнуть ситуация, когда нам необходимо, чтобы она была строго меньше. Для этого в LangChain существует RecursiveCharacterTextSplitter, который позволяет установить несколько символом, по которым будет происходить разделение. По умолчанию это ["\n\n", "\n", " ", ""] .

Сначала текст разделяется по первому символу из списка. Если длина полученного фрагмента слишком большая, то Splitter переходит к следующему символу и пытается разделить текст по нему.

Повторим предыдущий пример с использованием RecursiveCharacterTextSplitter:

from langchain_text_splitters import RecursiveCharacterTextSplitter
from transformers import AutoTokenizer

model_id = "microsoft/Phi-3-mini-4k-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)

text_splitter = RecursiveCharacterTextSplitter.from_huggingface_tokenizer(
    tokenizer, chunk_size=100, chunk_overlap=0
)

268 Primary school is divided into infant school (pupils from 5 to 7 years old) and junior school (from 8 to 11 years old). In some areas there are middle schools instead of junior schools, which take pupils from 9 to 12 years old.
Primary schools have from 50-200 pupils.

234 Secondary schools are usually much larger than primary schools and most children - over 80 percent - go to a comprehensive school at the age of 11. These schools are for all. Pupils do not need (to pass an exam to go to these schools.

В результате получим 12 текстов вместо 5 с длиной 111-410 символов. Это, конечно, снова больше максимальной длины, но результат стал заметно лучше.

HTMLHeaderTextSplitter

Такой разделитель текста анализирует структуру HTML документа и выделяет теги заголовков для создания логических разделов в тексте. В качестве аргументов необходимо передать:

Список кортежей заголовков, по которым будет происходить разделение. Например, ("h1", "Header 1"), где "h1" - тег заголовка, "Header 1" соответствующее ему значение(название).
return_each_element: bool. Если значение True, каждый фрагмент будет возвращаться как отдельный документ, по умолчанию False

Пример:

from langchain_text_splitters import HTMLHeaderTextSplitter

html_content = """
<html>
  <body>
    <h1>Introduction</h1>
    <p>Welcome to the introduction section.</p>
    
    <h2>Background</h2>
    <p>Some background details here.</p>
    
    <h3>Conclusion</h3>
    <p>Final thoughts.</p>
  </body>
</html>
"""

headers_to_split_on = [
    ("h1", "Header 1"),
    ("h2", "Header 2"),
    ("h3", "Header 3"),
]

splitter = HTMLHeaderTextSplitter(headers_to_split_on)
chunks = splitter.split_text(html_content)

for chunk in chunks:
    print(split)

Сейчас мы ожидаем увидеть 3 документа. Помимо содержимого текущего раздела мы также получаем metadata, которые содержат информацию о предыдущих заголовках.

page_content='Welcome to the introduction section.' metadata={'Header 1': 'Introduction'}
page_content='Some background details here.' metadata={'Header 1': 'Introduction', 'Header 2': 'Background'}
page_content='Final thoughts.' metadata={'Header 1': 'Introduction', 'Header 2': 'Background', 'Header 3': 'Conclusion'}

Пример 2. Теперь уберем один из заголовком и посмотрим на результат, документ остается прежним:

headers_to_split_on = [
    ("h1", "Header 1"),
    #("h2", "Header 2"), убираем тег h2
    ("h3", "Header 3"),
]

Результат:

page_content='Welcome to the introduction section.  
Some background details here.' metadata={'Header 1': 'Introduction'}
page_content='Final thoughts.' metadata={'Header 1': 'Introduction', 'Header 3': 'Conclusion'}

В этот раз получили всего 2 документа, заголовок h2 был проигнорирован, поэтому его содержимое было добавлено к содержимому заголовка h1

Пример 3. Результат с использованием return_each_element=True, все остальное остается прежним:

splitter = HTMLHeaderTextSplitter(headers_to_split_on, return_each_element=True)

Результат:

page_content='Welcome to the introduction section.' metadata={'Header 1': 'Introduction'}
page_content='Some background details here.' metadata={'Header 1': 'Introduction'}
page_content='Final thoughts.' metadata={'Header 1': 'Introduction', 'Header 3': 'Conclusion'}

Снова получили 3 документа, но два из них относятся к h1.

Помимо тегов h1, ... Можно использовать и другие, например, "div", "p".

JsonSplitter

JsonSplitter разделяет данные json формата. Сначала он просматривает глубину вложенности данных, после чего пытается разбить их на меньшие блоки. JsonSplitter также пытается сохранить целостность объектов, но она может быть нарушена, если будет неверно задан максимальный размер фрагмента.

Пример данных:

 {
    "id": 1,
    "name": "Leanne Graham",
    "username": "Bret",
    "email": "Sincere@april.biz",
    "address": {
      "street": "Kulas Light",
      "suite": "Apt. 556",
      "city": "Gwenborough",
      "zipcode": "92998-3874",
      "geo": {
        "lat": "-37.3159",
        "lng": "81.1496"
      }
    },
    "phone": "1-770-736-8031 x56442",
    "website": "hildegard.org",
    "company": {
      "name": "Romaguera-Crona",
      "catchPhrase": "Multi-layered client-server neural-net",
      "bs": "harness real-time e-markets"
    }

Попробуем разбить json документ, содержащий данные о 10 людях:

from langchain_text_splitters import RecursiveJsonSplitter

data = requests.get("https://jsonplaceholder.typicode.com/users").json()

splitter = RecursiveJsonSplitter(max_chunk_size=500)
json_chunks = splitter.split_json(data, convert_lists=True)

for chunk in json_chunks:
    print(chunk)

В результате 10 фрагментов с данными, целостность структуры сохранена.

{'0': {'id': 1, 'name': 'Leanne Graham', 'username': 'Bret', 'email': 'Sincere@april.biz', 'address': {'street': 'Kulas Light', 'suite': 'Apt. 556', 'city': 'Gwenborough', 'zipcode': '92998-3874', 'geo': {'lat': '-37.3159', 'lng': '81.1496'}}, 'phone': '1-770-736-8031 x56442', 'website': 'hildegard.org', 'company': {'name': 'Romaguera-Crona', 'catchPhrase': 'Multi-layered client-server neural-net', 'bs': 'harness real-time e-markets'}}}
{'1': {'id': 2, 'name': 'Ervin Howell', 'username': 'Antonette', 'email': 'Shanna@melissa.tv', 'address': {'street': 'Victor Plains', 'suite': 'Suite 879', 'city': 'Wisokyburgh', 'zipcode': '90566-7771', 'geo': {'lat': '-43.9509', 'lng': '-34.4618'}}, 'phone': '010-692-6593 x09125', 'website': 'anastasia.net', 'company': {'name': 'Deckow-Crist', 'catchPhrase': 'Proactive didactic contingency', 'bs': 'synergize scalable supply-chains'}}}
{'2': {'id': 3, 'name': 'Clementine Bauch', 'username': 'Samantha', 'email': 'Nathan@yesenia.net', 'address': {'street': 'Douglas Extension', 'suite': 'Suite 847', 'city': 'McKenziehaven', 'zipcode': '59590-4157', 'geo': {'lat': '-68.6102', 'lng': '-47.0653'}}, 'phone': '1-463-123-4447', 'website': 'ramiro.info', 'company': {'name': 'Romaguera-Jacobson', 'catchPhrase': 'Face to face bifurcated interface', 'bs': 'e-enable strategic applications'}}}
{'3': {'id': 4, 'name': 'Patricia Lebsack', 'username': 'Karianne', 'email': 'Julianne.OConner@kory.org', 'address': {'street': 'Hoeger Mall', 'suite': 'Apt. 692', 'city': 'South Elvis', 'zipcode': '53919-4257', 'geo': {'lat': '29.4572', 'lng': '-164.2990'}}, 'phone': '493-170-9623 x156', 'website': 'kale.biz', 'company': {'name': 'Robel-Corkery', 'catchPhrase': 'Multi-tiered zero tolerance productivity', 'bs': 'transition cutting-edge web services'}}}
{'4': {'id': 5, 'name': 'Chelsey Dietrich', 'username': 'Kamren', 'email': 'Lucio_Hettinger@annie.ca', 'address': {'street': 'Skiles Walks', 'suite': 'Suite 351', 'city': 'Roscoeview', 'zipcode': '33263', 'geo': {'lat': '-31.8129', 'lng': '62.5342'}}, 'phone': '(254)954-1289', 'website': 'demarco.info', 'company': {'name': 'Keebler LLC', 'catchPhrase': 'User-centric fault-tolerant solution', 'bs': 'revolutionize end-to-end systems'}}}
{'5': {'id': 6, 'name': 'Mrs. Dennis Schulist', 'username': 'Leopoldo_Corkery', 'email': 'Karley_Dach@jasper.info', 'address': {'street': 'Norberto Crossing', 'suite': 'Apt. 950', 'city': 'South Christy', 'zipcode': '23505-1337', 'geo': {'lat': '-71.4197', 'lng': '71.7478'}}, 'phone': '1-477-935-8478 x6430', 'website': 'ola.org', 'company': {'name': 'Considine-Lockman', 'catchPhrase': 'Synchronised bottom-line interface', 'bs': 'e-enable innovative applications'}}}
{'6': {'id': 7, 'name': 'Kurtis Weissnat', 'username': 'Elwyn.Skiles', 'email': 'Telly.Hoeger@billy.biz', 'address': {'street': 'Rex Trail', 'suite': 'Suite 280', 'city': 'Howemouth', 'zipcode': '58804-1099', 'geo': {'lat': '24.8918', 'lng': '21.8984'}}, 'phone': '210.067.6132', 'website': 'elvis.io', 'company': {'name': 'Johns Group', 'catchPhrase': 'Configurable multimedia task-force', 'bs': 'generate enterprise e-tailers'}}}
{'7': {'id': 8, 'name': 'Nicholas Runolfsdottir V', 'username': 'Maxime_Nienow', 'email': 'Sherwood@rosamond.me', 'address': {'street': 'Ellsworth Summit', 'suite': 'Suite 729', 'city': 'Aliyaview', 'zipcode': '45169', 'geo': {'lat': '-14.3990', 'lng': '-120.7677'}}, 'phone': '586.493.6943 x140', 'website': 'jacynthe.com', 'company': {'name': 'Abernathy Group', 'catchPhrase': 'Implemented secondary concept', 'bs': 'e-enable extensible e-tailers'}}}
{'8': {'id': 9, 'name': 'Glenna Reichert', 'username': 'Delphine', 'email': 'Chaim_McDermott@dana.io', 'address': {'street': 'Dayna Park', 'suite': 'Suite 449', 'city': 'Bartholomebury', 'zipcode': '76495-3109', 'geo': {'lat': '24.6463', 'lng': '-168.8889'}}, 'phone': '(775)976-6794 x41206', 'website': 'conrad.com', 'company': {'name': 'Yost and Sons', 'catchPhrase': 'Switchable contextually-based project', 'bs': 'aggregate real-time technologies'}}}
{'9': {'id': 10, 'name': 'Clementina DuBuque', 'username': 'Moriah.Stanton', 'email': 'Rey.Padberg@karina.biz', 'address': {'street': 'Kattie Turnpike', 'suite': 'Suite 198', 'city': 'Lebsackbury', 'zipcode': '31428-2261', 'geo': {'lat': '-38.2386', 'lng': '57.2232'}}, 'phone': '024-648-3804', 'website': 'ambrose.net', 'company': {'name': 'Hoeger LLC', 'catchPhrase': 'Centralized empowering task-force', 'bs': 'target end-to-end models'}}}

Но если бы мы неправильно указали максимальный размер фрагмента, например, поставили значение 100, то вместо 10 фрагментов получили бы 68.

2. VectorStores

После того как мы научились работать с текстом, было бы неплохо научиться его хранить и с ним взаимодействовать.

Векторные хранилища - специальные хранилища, которые помогают получать информацию на основе векторных представлений. Эти векторы называют Embeddings, они представляют собой числовое представление семантики данных, их часто используют для поиска неструктурированных данных.

LangChain поддерживает большое количество векторных хранилищ с различными реализациями поиска и предоставляет стандартный интерфейс для работы с ними.

При создание объекта хранилища необходимо передать embedding model, которая будет преобразовывать текст в векторное представление.

embedding model я возьму с HuggingFace, как делал это в первой части:

from langchain_huggingface.embeddings import HuggingFaceEndpointEmbeddings

embeddings = HuggingFaceEndpointEmbeddings(model="mixedbread-ai/mxbai-embed-large-v1", task="feature-extraction")

Для простых примеров воспользуемся InMemoryVectorStore, которая работает в оперативной памяти:

from langchain_core.vectorstores import InMemoryVectorStore

vectorstore = InMemoryVectorStore(embedding=embeddings)

Основные методы для работы:

add_documents - добавление списка текстов в хранилище.
delete_documents - удаление документов.
similarity_search - поиск документов похожих на заданный запрос.

Создадим несколько документов и добавим их. В качестве дополнительной информации возьмем источник:

from langchain_core.documents import Document

document_1 = Document(
    page_content="I had chocalate chip pancakes and scrambled eggs for breakfast this morning.",
    metadata={"source": "tweet"}
)

document_2 = Document(
    page_content="The weather forecast for tomorrow is cloudy and overcast, with a high of 62 degrees.",
    metadata={"source": "news"},
)

documents = [document_1, document_2]

#теперь можно просто добавить документы
vectorstore.add_documents(documents=documents)

#способ 2, дополнительно указать ID для документов, в будущем их можно будет удобно обновлять
vectorstore.add_documents(documents, ids=["doc1", "doc2"])

Теперь мы можем легко получить документ по его ID:

print(vectorstore.get_by_ids(["doc1"]))

Результат:

[Document(id='doc1', metadata={'source': 'tweet'}, page_content='I had chocalate chip pancakes and scrambled eggs for breakfast this morning.')]

Удаление документа:

vectorstore.delete(ids=["doc1"])

Поиск документов.

Векторные хранилища внедряют и сохраняют добавленные документы. Если мы передадим запрос, vectorstore внедрит запрос, выполнит поиск сходства по внедренным документам и вернет наиболее похожие. Это отражает две важные концепции: во-первых, должен быть способ измерить сходство между запросом и любым встроенным документом. Во-вторых, должен существовать алгоритм для эффективного выполнения поиска сходства по всем встроенным документам.

Некоторые способы сравнения векторов:

Косинусное подобие: измеряет косинус угла между двумя векторами.
Евклидово расстояние: измеряет расстояние по прямой между двумя точками.
Скалярное произведение: измеряет проекцию одного вектора на другой.

InMemoryVectorStore использует косинусное подобие.

Для поиска документа используется метод similarity_search, которому можно передать аргументы:

k - ограничение по количеству возвращаемых документов.
filter - позволяет отсортировать документы по metadata документа.

Попробуем найти документ по запросу. При этом поиск будем проводить только среди твитов.

Для этого напишем функцию для фильтрации:

def filter_function(doc: Document) -> bool:
    return doc.metadata.get("source") == "tweet"

и воспользуемся методом similarity_search:

print(vectorstore.similarity_search(
    "scrambled eggs for breakfast", #наш запрос
    k=2, #количество возвращаемых документов
    filter=filter_function #функция фильтрации
))

В результате получим:

[Document(id='doc1', metadata={'source': 'tweet'}, page_content='I had chocalate chip pancakes and scrambled eggs for breakfast this morning.')]

У нас был всего один документ, который относится к твитам, но если бы их было больше, то вернулись бы первые k документов. Документы расположены в порядке наибольшего соответствия.

Для большего контроля можно использовать метод similarity_search_with_score, который помимо документа возвращает его соответствие запросу.

results = vectorstore.similarity_search_with_score(
    query="scrambled eggs for breakfast",
    k=2,
)

for doc, score in results:
    print(score, doc.page_content)

Результат:

0.78 I had chocalate chip pancakes and scrambled eggs for breakfast this morning.
0.36 The weather forecast for tomorrow is cloudy and overcast, with a high of 62 degrees.

Как видим, первый документ намного больше подходит под наш запрос.

Помимо InMemoryVectorStore часто можно встретить такие хранилища, как FAISS и Chroma.

Мне нравится использовать Chroma, так как ее можно использовать без создание дополнительных аккаунтов и использовать локально в проекте.

Пример создания:

from langchain_huggingface.embeddings import HuggingFaceEndpointEmbeddings
from langchain_chroma import Chroma

embeddings = HuggingFaceEndpointEmbeddings(model="mixedbread-ai/mxbai-embed-large-v1", task="feature-extraction")

store = Chroma(
    collection_name="example_collection", #название базы
    embedding_function=embeddings,
    persist_directory="./chroma_db" #директория в которой будут размещаться фаилы
)

3. Retrievers

Как можно было заметить, я еще ни разу не использовал метод invoke, который является одним из главным методов всех элементов LangChain. Это связано с тем, что у обычных векторных хранилищ его нет, поэтому мы не может использовать хранилища напрямую в наших цепочках. Для таких целей существуют retrieval systems - поисковые системы, с помощью которых можно обращаться к VectorStore, графическим, реляционным базам и добавлять свою логику.

Перейти от VectorStore к Retriever довольно просто, возьмем для примера Chroma из примера выше и несколько документов, воспользуемся методом as_retriever, которому необходимо передать:

search_type - Определяет тип поиска, который должен выполнять поисковик. Например, "similarity" (по умолчанию), "mmr", "similarity_score_threshold".
search_kwargs - дополнительные параметры, например, k, score_threshold (минимальный уровень соответствия, если мы используем similarity_score_threshold), fetch_k (количество документов которые передаются в mmr алгоритм), filter.

retriever = store.as_retriever(
    search_type="mmr", search_kwargs={"k": 1}
)

MMR (Maximal Marginal Relevance) - это алгоритм, который используется для повышения разнообразия полученных документов при сохранении релевантности запросу пользователя.

Запрос к Retriever можно вызывать помощью invoke, в котором можно указать дополнительный аргумент filter:

print(retriever.invoke("scrambled eggs for breakfast", filter={"source": "tweet"}))

Результат:

[Document(metadata={'source': 'tweet'}, page_content='I had chocolate chip pancakes and scrambled eggs for breakfast this morning.')]

Custom Retrievers

Чтобы создать свой Retriever необходим класс, который будет наследником BaseRetriever и реализовать логику метода _get_relevant_documents. В этом методе мы можем описать любую логику, например, обращение к базам данных или к интернет ресурсам.

Пример создания своей простой поисковой системы:

from typing import List
from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain_core.documents import Document
from langchain_core.retrievers import BaseRetriever

class MyRetriever(BaseRetriever):
    """
    Будем возвращать k документов, которые содержат в себе
    запрос
    """
    documents: List[Document]
    k: int

    def _get_relevant_documents(
            self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> list[Document]:
        matching_documents = []

        for document in self.documents:
            if len(matching_documents) > self.k:
                return matching_documents

            if query.lower() in document.page_content.lower():
                matching_documents.append(document)
        return matching_documents

Добавим несколько документов:

documents = [
    Document(
        page_content="Dogs are great companions, known for their loyalty and friendliness.",
        metadata={"type": "dog", "trait": "loyalty"},
    ),
    Document(
        page_content="Cats are independent pets that often enjoy their own space.",
        metadata={"type": "cat", "trait": "independence"},
    ),
    Document(
        page_content="Rabbits are social animals that need plenty of space to hop around.",
        metadata={"type": "rabbit", "trait": "social"},
    ),
]

Найдем релевантные документы с помощью метода batch:

retriever = MyRetriever(documents=documents, k=3)
print(retriever.batch(["dog", "cat"]))

В результате получим все документы содержащие ключевые слова:

[[Document(metadata={'type': 'dog', 'trait': 'loyalty'}, page_content='Dogs are great companions, known for their loyalty and friendliness.')], [Document(metadata={'type': 'cat', 'trait': 'independence'}, page_content='Cats are independent pets that often enjoy their own space.')]]

Заключение

В этой статье я постарался разобрать разделение и хранение данных. Конечно, предварительная обработка на этом не заканчивается, так как в тексте могут встречаться специальные символы, например, смайлики, которые тоже нужно как то обработать.

Важно помнить, что обработка текста - один из главных этапов работы с языковыми моделями, так как от правильности и полноты входных данных зависит дальнейшая работа языковой модели.

В следующей части разберем создание собственных инструментов (Tools) и работу с ними.

Мой телеграмм канал, где я пишу на тему языковых моделей:

https://t.me/Viacheslav_Talks

Используем языковые модели в AI-агентах. Часть 2. Retrievers, TextSplitters