Upload 2 files

app.py

@@ -1,137 +1,141 @@
-import gradio as gr
-import pandas as pd
+from transformers import BertTokenizerFast, BertForSequenceClassification
+from sentence_transformers import SentenceTransformer
+from torch.nn.functional import cosine_similarity
 import torch
-import numpy as np
-from transformers import BertForSequenceClassification, BertTokenizerFast
 from datasets import load_dataset
-import ast  # Для преобразования строковых представлений списков
-
-# Убираем предупреждение о симлинках
-import os
-os.environ['HF_HUB_DISABLE_SYMLINKS_WARNING'] = '1'
+import pandas as pd
+import gradio as gr
 
-# Загрузка и подготовка данных
-def load_movie_data():
-    try:
-        dataset = load_dataset("wykonos/movies")
-        df = pd.DataFrame(dataset['train'])
-        
-        # Преобразование строковых представлений в списки
-        def parse_list(x):
-            try:
-                return [item['name'] for item in ast.literal_eval(x)]
-            except:
-                return []
-        
-        # Обработка колонок
-        df['genres'] = df['genres'].apply(parse_list)
-        df['production_companies'] = df['production_companies'].apply(parse_list)
-        df['keywords'] = df['keywords'].apply(parse_list)
-        
-        # Выбор нужных колонок
-        df = df[[
-            'id', 'title', 'genres', 'overview',
-            'release_date', 'vote_average', 'poster_path'
-        ]].drop_duplicates(subset=['id'])
-        
-        return df
+# Загрузка датасета
+def load_data():
+    dataset = load_dataset("wykonos/movies", split='train[:400]')
+    df = pd.DataFrame(dataset)
     
-    except Exception as e:
-        print(f"Ошибка загрузки данных: {str(e)}")
-        return pd.DataFrame()
-
-# Загрузка данных
-df = load_movie_data()
+    # Обработка жанров
+    df['genres'] = df['genres'].apply(
+        lambda x: x.split('-') if isinstance(x, str) else []
+    )
+    
+    # Исправление опечаток
+    genre_corrections = {
+        'hystery': 'mystery',
+        'adventura': 'adventure',
+        'nur': 'action'
+    }
+    
+    df['genres'] = df['genres'].apply(
+        lambda x: [genre_corrections.get(g.lower(), g).capitalize() for g in x]
+    )
+    
+    return df
 
-if df.empty:
-    raise RuntimeError("Не удалось загрузить данные. Проверьте структуру датасета.")
+df = load_data()
 
-# Загрузка модели и токенизатора
+# Загрузка модели
 model_name = "AventIQ-AI/bert-movie-recommendation-system"
-try:
-    tokenizer = BertTokenizerFast.from_pretrained(model_name)
-    model = BertForSequenceClassification.from_pretrained(model_name)
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = model.to(device)
-except Exception as e:
-    raise RuntimeError(f"Ошибка загрузки модели: {str(e)}")
-
-def predict_genres(text, threshold=0.5):
+tokenizer = BertTokenizerFast.from_pretrained(model_name)
+model = BertForSequenceClassification.from_pretrained(model_name)
+
+# Сопоставление меток модели с жанрами
+genre_labels = [
+    "Action", "Adventure", "Animation", "Comedy", "Crime",
+    "Documentary", "Drama", "Family", "Fantasy", "History",
+    "Horror", "Music", "Mystery", "Romance", "Science Fiction",
+    "TV Movie", "Thriller", "War", "Western"
+]
+
+model.config.id2label = {i: label for i, label in enumerate(genre_labels)}
+
+# Функция предсказания
+def predict_genres(text, threshold=0.3):
     inputs = tokenizer(
-        text, 
-        padding=True, 
-        truncation=True, 
-        max_length=128, 
+        text,
+        max_length=128,
+        padding=True,
+        truncation=True,
         return_tensors="pt"
-    ).to(device)
+    )
     
     with torch.no_grad():
         outputs = model(**inputs)
     
     probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
-    predicted_indices = np.where(probs.cpu().numpy() > threshold)[1]
+    predicted_indices = torch.where(probs > threshold)[1].tolist()
     
-    genre_labels = list(model.config.id2label.values())
-    predicted_genres = [genre_labels[i] for i in predicted_indices]
-    
-    return list(set(predicted_genres))
+    return list(set([model.config.id2label[i] for i in predicted_indices]))
+# Модель для текстовых эмбеддингов
+embedding_model = SentenceTransformer("paraphrase-multilingual-mpnet-base-v2")
+
+# Кодируем описания фильмов
+df["overview"] = df["overview"].fillna("")  # если есть пропуски
+vectors = embedding_model.encode(
+    df["overview"].tolist(),
+    convert_to_tensor=True,
+    show_progress_bar=True
+)
+
+df["overview_vector"] = list(vectors)
 
-def recommend_movies(query, top_k=10):
+
+def recommend_movies(query, top_k=5):
     try:
+        # 1. Предсказание жанров
         genres = predict_genres(query)
-        mask = df['genres'].apply(lambda x: any(g in genres for g in x))
-        filtered = df[mask]
         
-        # Сортировка по рейтингу и популярности
-        results = filtered.sort_values(
-            by=['vote_average', 'release_date'], 
-            ascending=[False, False]
-        ).head(top_k)
+        # 2. Кодируем сам запрос
+        query_vector = embedding_model.encode([query], convert_to_tensor=True)[0]
+
+        # 3. Считаем косинусное сходство с описаниями фильмов
+        similarities = cosine_similarity(
+            query_vector.unsqueeze(0),
+            torch.stack(df["overview_vector"].tolist())
+        ).squeeze(0)
+
+        df["similarity"] = similarities.cpu().numpy()
+
+        # 4. Объединяем семантику + жанры (опционально)
+        # Можно фильтровать по жанрам или просто повысить вес фильмам с совпадающими жанрами
+        def genre_score(row):
+            return any(g in row["genres"] for g in genres)
         
+        df["genre_boost"] = df.apply(genre_score, axis=1).astype(int)
+        df["total_score"] = df["similarity"] + df["genre_boost"] * 0.15  # вес жанра
+
+        # 5. Выбираем top_k
+        results = df.sort_values(
+            by="total_score", ascending=False
+        ).head(top_k)
+
+        # 6. Форматируем ответ
         output = []
         for _, row in results.iterrows():
-            movie_info = [
-                f"🎬 **{row['title']}** ({row['release_date'][:4]})",
+            info = [
+                f"🎬 {row['title']}",
                 f"⭐ Рейтинг: {row['vote_average']}",
-                f"📀 Жанры: {', '.join(row['genres'][:3])}",
+                f"🎭 Жанры: {', '.join(row['genres'])}",
+                f"📅 Год: {row['release_date'][:4] if pd.notna(row['release_date']) else 'N/A'}",
+                f"📖 {row['overview'][:300]}..."
             ]
-            
-            if row['overview']:
-                movie_info.append(f"📖 Описание: {row['overview'][:200]}...")
-                
-            if row['poster_path']:
-                movie_info.append(f"🖼️ Постер: https://image.tmdb.org/t/p/w500{row['poster_path']}")
-            
-            output.append("\n".join(movie_info) + "\n" + "-"*50)
+            output.append("\n".join(info))
         
-        return "\n\n".join(output) if len(output) > 0 else "Нет подходящих фильмов"
-    
+        return "\n\n" + "\n\n".join(output) if output else "Нет результатов"
+
     except Exception as e:
         return f"Ошибка: {str(e)}"
 
-# Создание интерфейса Gradio с HTML-стилями
-css = """
-.gradio-container {background: #f0f2f6}
-h1 {text-align: center; color: #2d3436}
-"""
 
+# Создание интерфейса
 interface = gr.Interface(
     fn=recommend_movies,
-    inputs=gr.Textbox(
-        label="🎥 Введите запрос", 
-        placeholder="Пример: Космическая опера с эпическими битвами..."
-    ),
-    outputs=gr.Markdown(label="🎬 Результаты поиска"),
-    title="🍿 AI-Кинотеатр: Персональные рекомендации",
-    description="Система рекомендаций фильмов на основе глубокого обучения",
+    inputs=gr.Textbox(label="Опишите желаемый фильм"),
+    outputs=gr.Textbox(label="Рекомендации"),
     examples=[
-        ["Комедия с неожиданным сюжетом"],
-        ["Драма о взрослении подростков"],
-        ["Фантастический боевик с роботами"]
+        ["Страшный фильм с привидениями"],
+        ["Веселая комедия про студентов"],
+        ["Фантастика с космическими битвами"]
     ],
-    css=css,
-    allow_flagging="never"
+    title="🍿 AI Киносоветник"
 )
 
 if __name__ == "__main__":
-    interface.launch()
+    interface.launch(server_name="localhost", server_port=7860)

requirements.txt

@@ -1,4 +1,5 @@
 sentence-transformers
 faiss-cpu
 gradio
-pandas
+pandas
+datasets