Помощник читателя: визуализируем сюжет

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

Откуда растут ноги.

Думаю, многие, кто окунается в любое, хоть сколько-нибудь сложное произведение, порою теряется в хитросплетениях взаимоотношений героев, причин их поступков и развитии общего настроения произведения. Особенно если вещают несколько рассказчиков, события подаются не в хронологическом порядке, имеет место реверсивная композиция, или линии развиваются параллельно. У хитрого-то писателя все расписано и всегда перед глазами. Кто есть кто, что у кого на уме, где случится встреча и когда выстрелит ружье. Для примера окопал фотографию своей шпаргалки, нарисованной в процессе первого прочтения «Бесов» Достоевского:

Читаем текст и генерируем эмбеддинги

В этом подходе я волей обстоятельств положился на YandexGPT, поэтому эмбеддинги и дальнейшие примеры приведены на основе LLM Яндекса.

Доступный в langchain_community класс YandexGPTEmbeddings дополнительно оборачиваем в лимитер, позволяющий не заспамить API, ограниченный десятью запросами в секунду (опускаю детали, полный код здесь):

from tenacity import (     retry,     stop_after_attempt,     wait_exponential,     retry_if_exception_type ) from pydantic import Field, ConfigDict, BaseModel from langchain_community.embeddings.yandex import YandexGPTEmbeddings # ...  class RateLimitedEmbeddings(YandexGPTEmbeddings):     # ...     @retry(         retry=retry_if_exception_type(Exception),         stop=stop_after_attempt(3),         wait=wait_exponential(multiplier=1, min=2, max=10)     )     def _embed_batch(self, batch: List[str]) -> List[List[float]]:         time.sleep(0.1)         return super().embed_documents(batch)      def embed_documents(self, texts: List[str]) -> List[List[float]]:         # ...         result = []                      for i in range(0, len(texts), self.batch_size):             batch = texts[i:i + self.batch_size]             # ...             batch_result = self._embed_batch(batch)             result.extend(batch_result)                          if i + self.batch_size < len(texts):             time.sleep(self.delay_between_batches)                      return result

Далее делаем следующее:

1. Посредством TextLoader читаем файлик с текстом

2. При помощи RecursiveCharacterTextSplitter разделяем текст на чанки, заданные параметрами chunk_size и chunk_overlap (здесь 1000 и 100 соответственно).

3. Генерируем эмбеддинги с помощью объявленного выше RateLimitedEmbeddings, и складываем их в векторное хранилище FAISS.

4. Инициализируем языковую модель YandexGPT (здесь yandexgpt-32k).

5. Создаём экземпляр RetrievalQA для ответа на вопросы по данным из векторного хранилища.

from langchain.document_loaders import TextLoader from langchain.text_splitter import RecursiveCharacterTextSplitter from langchain.vectorstores import FAISS from langchain.chains import RetrievalQA from langchain_community.llms import YandexGPT # ...  loader = TextLoader(file_path, encoding="utf-8") documents = loader.load()          text_splitter = RecursiveCharacterTextSplitter(     chunk_size=config.chunk_size,     chunk_overlap=config.chunk_overlap ) texts = text_splitter.split_documents(documents)  embeddings = RateLimitedEmbeddings()  vectorstore = FAISS.from_documents(texts, embeddings)          llm = YandexGPT(     api_key=config.api_key,     folder_id=config.folder_id,     model_uri=config.model_uri )  qa = RetrievalQA.from_chain_type(     llm=llm,     chain_type="map_reduce",     retriever=vectorstore.as_retriever(search_kwargs={"k": config.search_k}),     return_source_documents=False )

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

Граф связей между персонажами

qa(""" Представь связь всех главных героев книги в виде списка в формате JSON, где каждый элемент имеет формат:  {  "name": "имя одного героя",  "links": {  "имя другого героя": "тип отношений между ними",  ...  }  }  Тип отношений должен быть лаконичным, например: отец, сестра, друг, знакомый, cупруг и т.п. """)

[    {       "name":"Дмитрий Старцев (Ионыч)",       "links":{          "Иван Петрович Туркин":"знакомый",          "Вера Иосифовна Туркина":"пациентка",          "Екатерина Ивановна Туркина (Котик)":"объект любви"       }    },    {       "name":"Иван Петрович Туркин",       "links":{          "Дмитрий Старцев":"знакомый",          "Вера Иосифовна Туркина":"жена",          "Екатерина Ивановна Туркина":"дочь"       }    },    {       "name":"Вера Иосифовна Туркина",       "links":{          "Дмитрий Старцев":"пациент",          "Иван Петрович Туркин":"муж",          "Екатерина Ивановна Туркина":"дочь"       }    },    {       "name":"Екатерина Ивановна Туркина (Котик)",       "links":{          "Дмитрий Старцев":"объект симпатии",          "Иван Петрович Туркин":"отец",          "Вера Иосифовна Туркина":"мать"       }    } ]

Визуализируем полученную структуру дешево и сердито, с помощью matplotlib и networkx (весь код, опять же, здесь):

import networkx as nx import matplotlib.pyplot as plt # ...  G = nx.Graph()              for character in data:     name = character["name"]     G.add_node(name)     for linked_character, relation in character["links"].items():         G.add_edge(name, linked_character, relation=relation)  pos = nx.spring_layout(G) nx.draw(     G, pos,     with_labels=True,     node_color=self.config.graph.node_color,     node_size=self.config.graph.node_size,     font_size=self.config.graph.font_size )  edge_labels = nx.get_edge_attributes(G, "relation") nx.draw_networkx_edge_labels(     G, pos,     edge_labels=edge_labels,     font_size=self.config.graph.edge_font_size )              plt.title(self.config.graph.title) plt.axis("off") plt.show()

Для упомянутых вначале «Бесов» визуализатор выдает вот такой граф для главных героев произведения:

Изменение формулировки запроса к модели позволяет, например, описать характер отношений между героями, а не только формальный тип родственных связей. Для главы Бэла из «Героя нашего времени» получилась вот такая пентаграмма:

Очень важен размер контекста модели. Те же «Бесы» преобразуются в порядка 700k токенов. Такой размер способны объять лишь недавно появившиеся в публичном доступе модели.

Хронология событий

Движемся дальше — попросим модель представить хронологию событий в книге. Запрос выглядит следующим образом:

qa(""" Составь список событий в книге в формате JSON: [     {         "date": "Дата события по английски в английской локали",         "event": "Краткое описание события по русски"     } ] Ограничься только 10 событиями """)

Визуализируем по традиции максимально просто:

for event in data:     event['date'] = datetime.datetime.strptime(event['date'], '%d %B %y')  data.sort(key=lambda x: x['date'])  dates = [event['date'] for event in data] events = [event['event'] for event in data]  fig, ax = plt.subplots(figsize=self.config.timeline.figsize) ax.plot(     [1] * len(dates), dates,     marker='o',     color=self.config.timeline.marker_color,     linestyle=self.config.timeline.linestyle )  for i, event in enumerate(events):     ax.annotate(         event,         (1, dates[i]),         xytext=(10, 0),         textcoords='offset points',         ha='left',         va='center',         fontsize=self.config.timeline.fontsize     )  ax.yaxis.set_major_formatter(DateFormatter('%d %b %Y')) #...  plt.show()

Попросим модель нарисовать хронологию событий из дневника доктора Борменталя из Булгаковского «Собачьего сердца»:

Легко переваривать дневниковые записи. Тяжелее обстоит дело с хаотично разбросанными по тексту датами, особенно когда перемежаются собственно действие и какие-нибудь исторические справки. Хороший пример — Чеховский «Остров Сахалин». Если попросить модель составить список событий, произошедших с рассказчиком, выходит такая картина:

А если потребовать на стол список исторических событий, упомянутых в книге, то иная:

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

Карта мест действия

И напоследок менее тривиальная задача — нарисовать карту действий книги. Здесь между непосредственно инференсом модели и отрисовкой полученных данных добавляется этап геокодинга. Требуется получить географические координаты по полученным от модели топонимам.

Вначале просим модель составить список мест:

qa("""   Выведи список географических объектов из текста. Только названия через запятую. """)

Затем получаем их координаты:

from yandex_geocoder import Client # ...  locator = Client(YANDEX_GEOCODER_API_KEY)  locations = text.split(', ') result = []  for loc in set(locations):     # ...     coords = locator.coordinates(loc)     if coords:         result.append({"name": loc, "coordinates": [str(c) for c in coords]})

И уже теперь кладем их на карту с помощью cartopy,

import cartopy.crs as ccrs import cartopy.feature as cfeature # ...  longitudes = [float(coord[0]) for coord in [d['coordinates'] for d in data]] latitudes = [float(coord[1]) for coord in [d['coordinates'] for d in data]] names = [d['name'] for d in data] # ...  fig, ax = plt.subplots(     figsize=self.config.map.figsize,     subplot_kw={'projection': ccrs.PlateCarree()} )  ax.add_feature(cfeature.LAND) ax.add_feature(cfeature.OCEAN) ax.add_feature(cfeature.COASTLINE, linewidth=0.3) ax.add_feature(cfeature.BORDERS, linestyle=':', linewidth=0.3) ax.add_feature(cfeature.LAKES, alpha=0.5) ax.add_feature(cfeature.RIVERS)  for lon, lat, name in zip(longitudes, latitudes, names):     ax.plot(           lon, lat,           marker='o',           color=self.config.map.marker_color,           markersize=self.config.map.marker_size,           transform=ccrs.PlateCarree()     ) # ...      plt.show()

Пробуем что-то простенькое, например «Вокруг света за 80 дней»:

И что-то менее прямолинейное «На Западном фронте без перемен»:

В конечном счете все очень сильно упирается в размер модели. С локальными решениями добиться хорошего результата порою сложно. Не говоря уже о том, что в их контекстное окно с трудом влезает даже малая проза. А вот крупные SOTA модели, доступные по API, выдают отличный результат уже сейчас. Без труда генерируют стройный JSON по заданной простым промптом структуре, и не ошибаются в смысловой нагрузке.