Spaces:
Runtime error
Runtime error
import streamlit as st | |
import torch | |
import numpy as np | |
import pandas as pd | |
from PIL import Image | |
from transformers import AutoTokenizer, AutoModel | |
import re | |
import pickle | |
import requests | |
from io import BytesIO | |
st.title("Книжные рекомендации") | |
# Загрузка модели и токенизатора | |
model_name = "cointegrated/rubert-tiny2" | |
tokenizer = AutoTokenizer.from_pretrained(model_name) | |
model = AutoModel.from_pretrained(model_name, output_hidden_states=True) | |
# Загрузка датасета и аннотаций к книгам | |
books = pd.read_csv('all+++.csv') | |
books['author'].fillna('other', inplace=True) | |
annot = books['annotation'] | |
# Получение эмбеддингов аннотаций каждой книги в датасете | |
length = 256 | |
# Определение запроса пользователя | |
query = st.text_input("Введите запрос") | |
col1, col2 = st.columns(2) | |
generate_button = col2.button('Сгенерировать') | |
if generate_button: | |
with open("book_embeddings256xxx.pkl", "rb") as f: | |
book_embeddings = pickle.load(f) | |
query_tokens = tokenizer.encode_plus( | |
query, | |
add_special_tokens=True, | |
max_length=length, # Ограничение на максимальную длину входной последовательности | |
pad_to_max_length=True, # Дополним последовательность нулями до максимальной длины | |
return_tensors='pt' # Вернём тензоры PyTorch | |
) | |
with torch.no_grad(): | |
query_outputs = model(**query_tokens) | |
query_hidden_states = query_outputs.hidden_states[-1][:, 0, :] | |
query_hidden_states = torch.nn.functional.normalize(query_hidden_states) | |
# Вычисление косинусного расстояния между эмбеддингом запроса и каждой аннотацией | |
cosine_similarities = torch.nn.functional.cosine_similarity( | |
query_hidden_states.squeeze(0), | |
torch.stack(book_embeddings) | |
) | |
cosine_similarities = cosine_similarities.numpy() | |
indices = np.argsort(cosine_similarities)[::-1] # Сортировка по убыванию | |
num_books_per_page = col1.selectbox("Количество книг на странице:", [3, 5, 10], index=0) | |
for i in indices: | |
cols = st.columns(2) # Создание двух столбцов для размещения информации и изображения | |
cols[1].write("## " + books['title'][i]) | |
cols[1].markdown("**Автор:** " + books['author'][i]) | |
cols[1].markdown("**Аннотация:** " + books['annotation'][i]) | |
image_url = books['image_url'][i] | |
response = requests.get(image_url) | |
image = Image.open(BytesIO(response.content)) | |
cols[0].image(image) | |
cols[0].write(cosine_similarities[i]) | |
cols[1].write("---") | |
num_books_per_page -= 1 | |
if num_books_per_page == 0: | |
break | |