pages/model_w_clustering.py

import pandas as pd
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
import streamlit as st
import requests
from PIL import Image
from io import BytesIO
from langchain_community.chat_models.gigachat import GigaChat

st.title('Рекомендации сериалов по описанию пользователя с помощью симметричного семантического поиска')
st.divider()
df = pd.read_csv('clean_series_data.csv')
embeddings = np.load('embeddings.npy')


def load_image_from_url(url):
    try:
        response = requests.get(url)
        response.raise_for_status()
        return Image.open(BytesIO(response.content))
    except Exception as e:
        st.error(f"Не удалось загрузить изображение: {e}")
        return None


model = SentenceTransformer('cointegrated/rubert-tiny2')
model.cpu()
# embeddings_desc = df['Описание'].apply(lambda x: model.encode(x))
# embeddings_gan = df['Жанры'].apply(lambda x: model.encode(x))

# embeddings = embeddings_desc + embeddings_gan
metric = st.radio('Выберите метрику для поиска', [
                  'Евклидово расстояние', 'Косинусное сходство'])
if metric == 'Евклидово расстояние':
    embeddings = np.array(embeddings).astype(np.float32)
    faiss.normalize_L2(embeddings)

    dimension = embeddings.shape[1]

    nlist = 150

    quantizer = faiss.IndexFlatL2(dimension)

    index = faiss.IndexIVFFlat(quantizer, dimension, nlist, faiss.METRIC_L2)
    index.train(embeddings)
    # index = faiss.IndexFlatIP(dimension)

    index.add(embeddings)
    query = [st.text_area('Введите описание сериала')]

    button = st.button('Вывести результаты')
    if button:
        if query:
            query_embedding = model.encode(query).astype(np.float32)
            # Две строки ниже можно будет убрать
            query_embedding = np.array(
                query_embedding, dtype=np.float32).reshape(1, -1)
            faiss.normalize_L2(query_embedding)

            k = st.slider('Сколько сериалов рекомендовать?',
                          min_value=1, max_value=10, value=3, step=1)
            distances, indices = index.search(query_embedding, k)

            st.subheader('Похожие сериалы:')
            for i in range(k):
                url = df.loc[indices[0][i]]["Изображение"]
                image = load_image_from_url(url)
                st.image(image)
                st.write(f'Название: {df.loc[indices[0][i]]["Название"]}')
                st.write(f'Рейтинг: {df.loc[indices[0][i]]["Рейтинг"]}')
                st.write(f'Жанр: {df.loc[indices[0][i]]["Жанры"]}')
                st.write(f'Страна: {df.loc[indices[0][i]]["Страна"]}')
                st.write(
                    f'Длительность одной серии: {df.loc[indices[0][i]]["Длительность"]}')
                st.write(
                    f'Количество серий: {df.loc[indices[0][i]]["Количество серий"]}')
                st.write(f'Описание: {df.loc[indices[0][i]]["Описание"]}')
                st.write(f'Евклидово расстояние: {distances[0][i]:.4f}')
                st.divider()
else:
    embeddings = np.array(embeddings).astype(np.float32)
    faiss.normalize_L2(embeddings)

    dimension = embeddings.shape[1]
    nlist = 150
    quantizer = faiss.IndexFlatIP(dimension)
    index = faiss.IndexIVFFlat(
        quantizer, dimension, nlist, faiss.METRIC_INNER_PRODUCT)
    index.train(embeddings)
    # index = faiss.IndexFlatIP(dimension)

    index.add(embeddings)

    query = [st.text_area('Введите описание сериала')]

    button = st.button('Вывести результаты')
    if button:
        if query:
            query_embedding = model.encode(query).astype(np.float32)
            # Две строки ниже можно будет убрать
            query_embedding = np.array(
                query_embedding, dtype=np.float32).reshape(1, -1)
            faiss.normalize_L2(query_embedding)

            k = st.slider('Сколько сериалов рекомендовать?',
                          min_value=1, max_value=10, value=3, step=1)
            distances, indices = index.search(query_embedding, k)

            st.subheader('Похожие сериалы:')
            for i in range(k):
                url = df.loc[indices[0][i]]["Изображение"]
                image = load_image_from_url(url)
                st.image(image)
                st.write(f'Название: {df.loc[indices[0][i]]["Название"]}')
                st.write(f'Рейтинг: {df.loc[indices[0][i]]["Рейтинг"]}')
                st.write(f'Жанр: {df.loc[indices[0][i]]["Жанры"]}')
                st.write(f'Страна: {df.loc[indices[0][i]]["Страна"]}')
                st.write(
                    f'Длительность одной серии: {df.loc[indices[0][i]]["Длительность"]}')
                st.write(
                    f'Количество серий: {df.loc[indices[0][i]]["Количество серий"]}')
                st.write(f'Описание: {df.loc[indices[0][i]]["Описание"]}')
                st.write(f'Косинусное сходство: {distances[0][i]:.4f}')
                st.divider()

st.subheader('Генерация краткого содержания сериала с помощью SberGigaChat')

name_of_series = st.text_input('Введите название сериала')
gen_button = st.button('Показать краткое содержание')
giga = GigaChat(
    credentials='MjA2MGEzNjItZjE0Mi00NWE5LTllMDItMWVjZWRlNDA2ODM0OjNhNzNlZDJmLTY4NWUtNDI1Zi1iZjg4LTkxOWFjMjkxZDg0OA==', verify_ssl_certs=False)
if gen_button:
    with st.spinner('Генерация текста...'):
        st.write(giga.invoke(
            f"Расскажи cюжет сериала {name_of_series}").content)
st.divider()