File size: 5,209 Bytes
dae6371 27090a6 dae6371 27090a6 dae6371 27090a6 dae6371 27090a6 dae6371 27090a6 dae6371 27090a6 dae6371 27090a6 dae6371 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 |
import gradio as gr
from sentence_transformers import SentenceTransformer, util
import json
import os
# Загружаем модель
model_name = "HIT-TMG/KaLM-embedding-multilingual-mini-instruct-v1"
model = SentenceTransformer(model_name)
# Имя файла для сохранения эмбеддингов
embeddings_file = f"movie_embeddings_{model_name.replace('/', '_')}.json"
# Загружаем данные из файла movies.json
try:
with open("movies.json", "r", encoding="utf-8") as f:
movies_data = json.load(f)
except FileNotFoundError:
print("Ошибка: Файл movies.json не найден.")
movies_data = []
# Проверяем, есть ли сохраненные эмбеддинги
if os.path.exists(embeddings_file):
with open(embeddings_file, "r", encoding="utf-8") as f:
movie_embeddings_loaded = json.load(f)
print("Загружены эмбеддинги из файла.")
else:
movie_embeddings_loaded = {}
# Создаем словарь с описаниями фильмов и эмбеддингами
movie_descriptions = {}
movie_embeddings = {}
for movie in movies_data:
title = movie["name"]
# Формируем строку для эмбеддинга из всех данных фильма
embedding_string = f"Название: {movie['name']}\nГод: {movie['year']}\nЖанры: {movie['genresList']}\nОписание: {movie['description']}"
movie_descriptions[title] = embedding_string
# Проверяем, есть ли эмбеддинг для фильма в загруженных
if title in movie_embeddings_loaded:
movie_embeddings[title] = movie_embeddings_loaded[title]
else:
# Создаем эмбеддинг и добавляем в словарь
embedding = model.encode(embedding_string, convert_to_tensor=True).tolist()
movie_embeddings[title] = embedding
# Сохраняем эмбеддинги в файл, если были созданы новые
if len(movie_embeddings_loaded) < len(movie_embeddings):
with open(embeddings_file, "w", encoding="utf-8") as f:
json.dump(movie_embeddings, f, ensure_ascii=False, indent=4)
print("Эмбеддинги сохранены в файл.")
# Конвертируем эмбеддинги в тензоры
if len(movie_embeddings) > 0:
movie_embeddings_tensor = {
title: util.pytorch_cos_sim(
model.encode(query, convert_to_tensor=True),
model.encode(embedding_string, convert_to_tensor=True)
)
for title, embedding_string in movie_descriptions.items()
}
else:
movie_embeddings_tensor = None
def search_movies(query, top_k=3):
"""
Ищет наиболее похожие фильмы по запросу.
Args:
query: Текстовый запрос.
top_k: Количество возвращаемых результатов.
Returns:
Строку с результатами поиска в формате HTML.
"""
if movie_embeddings_tensor is None:
return "<p>Ошибка: Данные фильмов не загружены.</p>"
# Сортируем фильмы по убыванию сходства с запросом
sorted_movies = sorted(
movie_embeddings_tensor.items(),
key=lambda item: util.pytorch_cos_sim(
model.encode(query, convert_to_tensor=True),
model.encode(movie_descriptions[item[0]], convert_to_tensor=True)
)[0][0],
reverse=True
)
results_html = ""
for title, _ in sorted_movies[:top_k]:
# Ищем полное описание фильма в исходных данных
for movie in movies_data:
if movie["name"] == title:
description = movie["description"]
year = movie["year"]
genres = movie["genresList"]
score = util.pytorch_cos_sim(
model.encode(query, convert_to_tensor=True),
model.encode(movie_descriptions[title], convert_to_tensor=True)
)[0][0].item()
break
results_html += f"<h3><b>{title} ({year})</b></h3>"
results_html += f"<p><b>Жанры:</b> {genres}</p>"
results_html += f"<p><b>Описание:</b> {description}</p>"
results_html += f"<p><b>Сходство:</b> {score:.4f}</p>"
results_html += "<hr>"
return results_html
# Создаем интерфейс Gradio
iface = gr.Interface(
fn=search_movies,
inputs=gr.Textbox(label="Введите запрос:"),
outputs=gr.HTML(label="Результаты поиска:"),
title="Поиск фильмов по описанию",
description="Введите запрос, и система найдет наиболее похожие фильмы по их описаниям.",
examples=[
["Фильм про ограбление"],
["Комедия 2019 года"],
["Фантастика про космос"],
],
)
# Запускаем приложение
iface.launch() |