Update app.py

app.py

@@ -1,25 +1,66 @@
 import gradio as gr
 from sentence_transformers import SentenceTransformer, util
+import json
+import os
 
 # Загружаем модель
-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)
+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):
     """
@@ -32,19 +73,35 @@ def search_movies(query, top_k=3):
     Returns:
         Строку с результатами поиска в формате HTML.
     """
-    query_embedding = model.encode(query, convert_to_tensor=True)
-    cosine_scores = util.cos_sim(query_embedding, movie_embeddings)
+    if movie_embeddings_tensor is None:
+        return "<p>Ошибка: Данные фильмов не загружены.</p>"
 
-    # Получаем индексы топ-k наиболее похожих фильмов
-    top_results = cosine_scores.argsort(descending=True)[0][:top_k]
+    # Сортируем фильмы по убыванию сходства с запросом
+    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 i in top_results:
-        title = list(movie_descriptions.keys())[i]
-        description = list(movie_descriptions.values())[i]
-        score = cosine_scores[0][i].item()
+    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}</b></h3>"
+        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>"
@@ -59,9 +116,9 @@ iface = gr.Interface(
     title="Поиск фильмов по описанию",
     description="Введите запрос, и система найдет наиболее похожие фильмы по их описаниям.",
     examples=[
-        ["Фильм про мальчика волшебника и борьбу со злом"],
-        ["Хочу посмотреть фильм про путешествие, чтобы уничтожить артефакт"],
-        ["Фантастика про поиск новой планеты"],
+        ["Фильм про ограбление"],
+        ["Комедия 2019 года"],
+        ["Фантастика про космос"],
     ],
 )