from transformers import BertTokenizerFast, BertModel
from sklearn.metrics.pairwise import cosine_similarity
import pandas as pd
import gradio as gr

tokenizer_bert = BertTokenizerFast.from_pretrained("kykim/bert-kor-base")
model_bert = BertModel.from_pretrained("kykim/bert-kor-base")

df = pd.read_pickle('BookData_real_real_final.pkl')
df_emb = pd.read_pickle('review_emb.pkl')

title = "🍀고민 해결 도서 추천🍀"
description = "당신의 고민 해결을 도와줄 책을 추천 해드립니다♥"
examples = [["요즘 잠이 안 와"], ["앞으로 뭘 해야 할까?"]]

def embed_text(text):
    inputs = tokenizer_bert(text, return_tensors="pt")
    outputs = model_bert(**inputs)
    embeddings = outputs.last_hidden_state.mean(dim=1)  # 평균 임베딩 사용
    return embeddings.detach().numpy()[0]

def recommend(message):
  
  columns = ['거리']
  list_df = pd.DataFrame(columns=columns)

  emb = embed_text(message)
  list_df['거리'] = df_emb['서평임베딩'].map(lambda x: cosine_similarity([emb], [x]).squeeze())
  answer = df.loc[list_df['거리'].idxmax()]
  book_title = answer['제목']
  book_author = answer['작가']
  book_publi = answer['출판사']
  return "[" + book_author + "] 작가님의 「" + book_title + "」 추천합니다😊" + " (출판사:" + book_publi + ")"


iface = gr.Interface(fn=recommend,
                     inputs="text",
                     outputs="text",
                     theme="finlaymacklon/boxy_violet",
                     title=title,
                     description=description,
                     examples=examples)
iface.launch()