main.py

import numpy as np
import faiss
import streamlit as st
from transformers import AutoTokenizer, AutoModel
import torch
import joblib
import pandas as pd

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

# Загрузка сохраненных данных и индекса
text_embeddings = joblib.load('mail_embeddings.joblib')
index = faiss.read_index('mail_faiss_index.index')

# Датасет
df = pd.read_csv('clean_mail_movie.csv')
titles = df['movie_title'].tolist()
images = df['image_url'].tolist()
descr = df['description'].tolist()
links = df['page_url'].tolist()

# Загрузка модели и токенизатора
tokenizer = AutoTokenizer.from_pretrained("cointegrated/rubert-tiny2")
model = AutoModel.from_pretrained("cointegrated/rubert-tiny2")

# Функция для векторизации текста
def embed_bert_cls(text, model, tokenizer):
    t = tokenizer(text, padding=True, truncation=True, return_tensors='pt', max_length=1024)
    with torch.no_grad():
        model_output = model(**{k: v.to(model.device) for k, v in t.items()})
    embeddings = model_output.last_hidden_state[:, 0, :]
    embeddings = torch.nn.functional.normalize(embeddings)
    return embeddings[0].cpu().numpy()



# Streamlit интерфейс
st.title("Умный поиск фильмов")

user_input = st.text_area("Введите описание фильма:")
num_recs = st.selectbox("Количество рекомендаций:", [1, 3, 5, 10])

if st.button("Найти"):
    if user_input:
        user_embedding = embed_bert_cls(user_input, model, tokenizer).astype('float32').reshape(1, -1)
        distances, top_indices = index.search(user_embedding, num_recs)  # Здесь добавляем переменную distances
        
        st.write(f"Рекомендованные фильмы (Топ-{num_recs}):")
        
        for i, index in enumerate(top_indices[0]):
            col1, col2, col3 = st.columns([1, 4, 1])  # Добавляем ещё одну колонку для уверенности
            
            with col1:
                try:
                    st.image(images[index])  # Загружаем обложку фильма
                except Exception as e:
                    st.write(f"Could not display image at index {index}. Error: {e}")  # Это на случай отсутствия обложки

            with col2:
                st.markdown(f"[{titles[index]}]({links[index]})")  # Название фильма сделано кликабельным
                st.write(descr[index])  # Выводим описание фильма
            
            with col3:
                st.write(f"Уверенность: {(1-sigmoid(distances[0][i]))*100:.1f}%")  # Выводим уверенность