В текущих кодогенеративных реалиях создать что-то новое с нуля до уровня худо-бедной демонстрации стало предательски просто. Только успевай доходчиво формулировать свои хотелки, да вовремя давать по рукам бездушной 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Далее делаем следующее:
Посредством TextLoader читаем файлик с текстом
При помощи
RecursiveCharacterTextSplitterразделяем текст на чанки, заданные параметрамиchunk_sizeиchunk_overlap(здесь 1000 и 100 соответственно).Генерируем эмбеддинги с помощью объявленного выше
RateLimitedEmbeddings, и складываем их в векторное хранилище FAISS.Инициализируем языковую модель YandexGPT (здесь
yandexgpt-32k).Создаём экземпляр
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упруг и т.п.
""")На выходе получаем желаемый JSON, пример для Чеховского Ионыча:
[
{
"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 по заданной простым промптом структуре, и не ошибаются в смысловой нагрузке.
