|
import gradio as gr |
|
from sentence_transformers import SentenceTransformer, util |
|
|
|
|
|
model = SentenceTransformer("HIT-TMG/KaLM-embedding-multilingual-mini-instruct-v1") |
|
|
|
|
|
movie_descriptions = { |
|
"Гарри Поттер и Философский камень": "Юный волшебник Гарри Поттер поступает в школу чародейства и волшебства Хогвартс, где ему предстоит не только обучиться магии, но и сразиться с темным волшебником Волан-де-Мортом, который убил его родителей и оставил шрам на лбу Гарри.", |
|
"Властелин колец: Братство Кольца": "Могущественное Кольцо Всевластия, созданное темным властелином Сауроном, попадает к хоббиту Фродо Бэггинсу. Чтобы не допустить порабощения Средиземья, Фродо и его друзья отправляются в опасное путешествие к Роковой Горе, чтобы уничтожить Кольцо в её пламени.", |
|
"Голодные игры": "В постапокалиптическом мире государство Панем ежегодно устраивает жестокие игры на выживание, в которых участвуют по одному юноше и девушке из каждого округа. Китнисс Эвердин добровольно вызывается участвовать вместо своей младшей сестры и борется за выживание в смертельной схватке.", |
|
"Семь": "Опытный детектив Уильям Сомерсет и его молодой напарник Дэвид Миллз расследуют серию загадочных убийств, совершенных в соответствии с семью смертными грехами. Каждый новый труп — это послание от изощренного убийцы, который возомнил себя орудием Божьего суда.", |
|
"Каникулы": "Семейство Грисвольдов отправляется в долгожданное путешествие на автомобиле через всю страну, мечтая провести незабываемый отпуск. Однако, с самого начала всё идет не по плану, и их поездка превращается в череду комичных и порой опасных приключений.", |
|
"Терминатор 2: Судный день": "Из постапокалиптического будущего прибывают два киборга: усовершенствованная модель Т-1000, способная принимать любую форму, и перепрограммированный Т-800. Цель первого — убить Джона Коннора, будущего лидера сопротивления, цель второго — защитить его любой ценой.", |
|
"Интерстеллар": "Когда засуха и пылевые бури ставят человечество на грань вымирания, группа исследователей отправляется через червоточину в другую галактику, чтобы найти новую планету, пригодную для жизни.", |
|
"Парк юрского периода": "Группа ученых прибывает на отдаленный остров, где эксцентричный миллиардер Джон Хэммонд создал парк развлечений с живыми динозаврами, клонированными из древней ДНК. Во время экскурсии система безопасности выходит из строя, и динозавры оказываются на свободе.", |
|
"Одержимость": "Молодой и амбициозный барабанщик Эндрю Ниман мечтает стать великим джазовым музыкантом. Чтобы достичь своей цели, он поступает в престижную консерваторию, где сталкивается с безжалостным дирижером Теренсом Флетчером, который использует жестокие методы обучения.", |
|
"Титаник": "В 1912 году роскошный лайнер «Титаник» отправляется в свое первое плавание через Атлантический океан. На борту встречаются Джек Доусон, бедный художник, и Роза Дьюитт Бьюкейтер, аристократка, помолвленная с богатым, но жестоким Кэлом Хокли. Между ними вспыхивает страстная любовь, которой предстоит пройти испытание катастрофой." |
|
} |
|
|
|
|
|
movie_embeddings = model.encode(list(movie_descriptions.values()), convert_to_tensor=True) |
|
|
|
def search_movies(query, top_k=3): |
|
""" |
|
Ищет наиболее похожие фильмы по запросу. |
|
|
|
Args: |
|
query: Текстовый запрос. |
|
top_k: Количество возвращаемых результатов. |
|
|
|
Returns: |
|
Строку с результатами поиска в формате HTML. |
|
""" |
|
query_embedding = model.encode(query, convert_to_tensor=True) |
|
cosine_scores = util.cos_sim(query_embedding, movie_embeddings) |
|
|
|
|
|
top_results = cosine_scores.argsort(descending=True)[0][:top_k] |
|
|
|
results_html = "" |
|
for i in top_results: |
|
title = list(movie_descriptions.keys())[i] |
|
description = list(movie_descriptions.values())[i] |
|
score = cosine_scores[0][i].item() |
|
|
|
results_html += f"<h3><b>{title}</b></h3>" |
|
results_html += f"<p><b>Описание:</b> {description}</p>" |
|
results_html += f"<p><b>Сходство:</b> {score:.4f}</p>" |
|
results_html += "<hr>" |
|
|
|
return results_html |
|
|
|
|
|
iface = gr.Interface( |
|
fn=search_movies, |
|
inputs=gr.Textbox(label="Введите запрос:"), |
|
outputs=gr.HTML(label="Результаты поиска:"), |
|
title="Поиск фильмов по описанию", |
|
description="Введите запрос, и система найдет наиболее похожие фильмы по их описаниям.", |
|
examples=[ |
|
["Фильм про мальчика волшебника и борьбу со злом"], |
|
["Хочу посмотреть фильм про путешествие, чтобы уничтожить артефакт"], |
|
["Фантастика про поиск новой планеты"], |
|
], |
|
) |
|
|
|
|
|
iface.launch() |
