Hugging Face's logo

Spaces:
SaHa24D
/
testsaha24
Sleeping

App Files Community
testsaha24
File size: 4,941 Bytes 
8266711
 
 
 
 
 
 
 
 
 
a52f0a3
 
8266711
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8cd0d6b
8266711
8cd0d6b
8266711
 
 
 
 
 
f37daba
8266711
f37daba
 
 
 
 
 
 
 
8266711
 
93f4c22
 
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
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
 
from transformers import BertTokenizerFast, BertForSequenceClassification
from sentence_transformers import SentenceTransformer
from torch.nn.functional import cosine_similarity
import torch
from datasets import load_dataset
import pandas as pd
import gradio as gr

# Загрузка датасета
def load_data():
    dataset = load_dataset("wykonos/movies", split='train[:1400]')
    df = pd.DataFrame(dataset)
    
    # Обработка жанров
    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

df = load_data()

# Загрузка модели
model_name = "AventIQ-AI/bert-movie-recommendation-system"
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,
        max_length=128,
        padding=True,
        truncation=True,
        return_tensors="pt"
    )
    
    with torch.no_grad():
        outputs = model(**inputs)
    
    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
    predicted_indices = torch.where(probs > threshold)[1].tolist()
    
    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=5):
    try:
        # 1. Предсказание жанров
        genres = predict_genres(query)
        
        # 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():
            info = [
                f"🎬 {row['title']}",
                f"⭐ Рейтинг: {row['vote_average']}",
                f"🎭 Жанры: {', '.join(row['genres'])}",
                f"📅 Год: {row['release_date'][:4] if pd.notna(row['release_date']) else 'N/A'}",
                f"📖 {row['overview'][:300]}..."
            ]
            output.append("\n".join(info))
        
        return "\n\n".join(output) if output else "Нет результатов"

    except Exception as e:
        return f"Ошибка: {str(e)}"


# Создание интерфейса
interface = gr.Interface(
    fn=recommend_movies,
    inputs=gr.Textbox(label="Опишите желаемый фильм"),
    outputs=gr.Textbox(label="Рекомендации"),
    examples=[
        ["Страшный фильм с привидениями"],
        ["Веселая комедия про студентов"],
        ["Фантастика с космическими битвами"]
    ],
    title="🍿 AI Киносоветник"
)
if __name__ == "__main__":
    interface.launch()