Spaces:
Sleeping
Sleeping
# Основные библиотеки | |
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("---") |