from transformers import BertTokenizerFast, BertModel
import gradio as gr

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

from sklearn.metrics.pairwise import cosine_similarity

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


def embed_text(text, tokenizer=tokenizer_bert, model=model_bert):
    inputs = tokenizer(text, return_tensors="pt")
    outputs = model(**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['제목']

  return book_title





title = "🍀고민 해결 도서 추천 챗봇🍀"
description = "당신의 고민 해결을 도와줄 책을 추천 해드립니다"
examples = [["요즘 잠이 안 와"]]




gr.ChatInterface(
    fn=recommend,
    title=title,
    description=description,
    examples=examples,
    inputs=["text", "state"],
    outputs=["chatbot", "state"],
    theme="finlaymacklon/boxy_violet",
).launch()