test / app.py
AlexandraGulamova's picture
test
343af91
# Основные библиотеки
import os
import re
import string
import warnings
import numpy as np
import pandas as pd
import torch
# Машинное обучение и обработка текста
from transformers import AutoTokenizer, AutoModel, AutoModelForSeq2SeqLM, MarianMTModel, MarianTokenizer
from sentence_transformers import SentenceTransformer, util
from sklearn.base import BaseEstimator, TransformerMixin
# FAISS для семантического поиска
import faiss
# Лемматизация и морфология
import pymorphy3
# Streamlit для создания веб-приложений
import streamlit as st
# Кастомные модули
from EmbeddingGenerator import EmbeddingGenerator
from TextAugmentation import TextAugmentation
# LangChain для интеграции GigaChat
from langchain_community.chat_models.gigachat import GigaChat
# ======= загружаем ранее рассчитанные эмбеддинги и объявляем классы=======
# Инициализация GigaChat с ключом и отключенной проверкой SSL
giga = GigaChat(
credentials="ODk0NDE1ODEtYTJhMi00N2Y1LTk4YWItNGZlNzNkM2QwMDNiOjk5YmVjN2ZjLThmM2EtNDhjYy04OWQ2LWNkOTlhOTNkNGY3NQ==",
verify_ssl_certs=False
)
augmentor = TextAugmentation()
embedding_gen = EmbeddingGenerator()
df=pd.read_csv("movies_data_fixed.csv")
image_path = "image-2.png"
# Загружаем и отображаем картинку
st.image(image_path, use_container_width=True)
# Заголовок
st.markdown(
"""
<div class="title">
КиноКринж
</div>
""",
unsafe_allow_html=True
)
# Добавляем окно ввода текста
user_input = st.text_area("Добавьте описание фильма", "", height=150)
# Слайдер для выбора количества фильмов
num_results = st.slider('Выберите количество фильмов', min_value=1, max_value=20, value=4)
# Выбор модели
model_option = st.selectbox('Выберите модель для обработки запроса:', ['cointegrated/rubert-tiny2','DeepPavlov/rubert-base-cased','all-MiniLM-L6-v2', 'paraphrase-MiniLM-L6-v2'])
if model_option!='DeepPavlov/rubert-base-cased':
model = SentenceTransformer(model_option)
# ======= дополнительная фильтрация для аугментаций (убираем слишком непохожие) =======
def filter_paraphrases(original, paraphrases, threshold=0.8):
original_embedding = model.encode(original)
filtered = []
for paraphrase in paraphrases:
paraphrase_embedding = model.encode(paraphrase)
similarity = util.cos_sim(original_embedding, paraphrase_embedding).item()
if similarity >= threshold:
filtered.append(paraphrase)
return filtered
#======================СЕМПЛ======= =======
# Проверка наличия рекомендованных фильмов
if 'recommended_movies' not in st.session_state:
st.session_state.recommended_movies = []
# Кнопка для поиска
if st.button('Найти фильм'):
if user_input.strip():
# Генерация эмбеддинга для запроса
if model_option != 'DeepPavlov/rubert-base-cased' and model_option != 'cointegrated/rubert-tiny2':
index = faiss.read_index('faiss_index.bin')
query_embedding = model.encode([user_input]).astype("float32")
faiss.normalize_L2(query_embedding)
elif model_option == 'DeepPavlov/rubert-base-cased':
index = faiss.read_index('pavlov3.bin')
back_translate = augmentor.back_translate(user_input)
augmented_query_pavlov = user_input + " " + back_translate
query_embedding = embedding_gen.generate_embeddings(augmented_query_pavlov, method="pavlov")
elif model_option == 'cointegrated/rubert-tiny2':
index = faiss.read_index('rubert2.bin')
paraphrase = augmentor.paraphrase(user_input, num_return_sequences=3)
filtered_rubert = filter_paraphrases(user_input, paraphrase)
augmented_query_rubert = user_input + " " + " ".join(filtered_rubert)
query_embedding = embedding_gen.generate_embeddings(augmented_query_rubert, method="rubert_tiny2").reshape(1, -1)
faiss.normalize_L2(query_embedding)
# Поиск ближайших соседей
distances, indices = index.search(query_embedding, num_results)
# Отображение результатов
st.write(f"Результаты поиска ({num_results} фильмов):")
recommended_movies = []
for idx, distance in zip(indices[0], distances[0]):
recommended_movies.append({
'title': df.iloc[idx]['movie_title'],
'description': df.iloc[idx]['description'],
'image_url': df.iloc[idx]['image_url'],
'page_url': df.iloc[idx]['page_url'],
'similarity': distance,
'short_description': None, # Содержимое краткого описания
'is_short_description_shown': False # Флаг для того, чтобы избежать повторного запроса
})
# Сохраняем результаты в session_state
st.session_state.recommended_movies = recommended_movies
# Отображение рекомендованных фильмов
for idx, movie in enumerate(st.session_state.recommended_movies):
st.write(f"### {movie['title']}")
st.write(f"Описание: {movie['description']}")
st.write(f"Схожесть: {movie['similarity']:.4f}")
# Отображаем картинку постера
if movie.get('image_url'):
st.image(movie['image_url'], width=200)
# Добавляем ссылку на страницу фильма
if movie.get('page_url'):
st.markdown(f"[Перейти на страницу фильма]({movie['page_url']})")
# Генерируем уникальный ключ с использованием индекса
button_key = f"short_description_button_{idx}" # Уникальный ключ для кнопки
if st.button(f"Получить краткое содержание для {movie['title']}", key=button_key):
if not movie.get('is_short_description_shown', False): # Проверяем состояние
try:
# Отправляем запрос в GigaChat
prompt = f"{movie['title']} краткое содержание фильма не более 100 слов"
response = giga.invoke(prompt)
# Извлекаем описание из ответа
description = response.content if response else "Описание не найдено."
movie['short_description'] = description
movie['is_short_description_shown'] = True
except Exception as e:
st.error(f"Ошибка при запросе в GigaChat: {e}")
# Показываем краткое содержание
if movie.get('short_description') and movie.get('is_short_description_shown', False):
st.write(f"Краткое содержание для {movie['title']}: {movie['short_description']}")
st.write("---")