import streamlit as st
from transformers import AutoTokenizer, AutoModel
import torch

# carregar o modelo e o tokenizer
model = AutoModel.from_pretrained("neuralmind/bert-large-portuguese-cased")
tokenizer = AutoTokenizer.from_pretrained("neuralmind/bert-large-portuguese-cased")

st.title("Comparador de Similaridade de Textos em Português")
st.write("Este aplicativo usa o modelo BERT-large treinado para o português para calcular a similaridade entre dois textos.")

# definir as caixas de entrada para as duas frases
frase1 = st.text_input("Insira a primeira frase:")
frase2 = st.text_input("Insira a segunda frase:")

# verificar se as frases foram inseridas e calcular a similaridade se sim
if frase1 and frase2:
    # tokenizar as frases e obter os embeddings
    input_ids = tokenizer.encode_plus(frase1, frase2, padding='max_length', truncation=True, return_tensors='pt')

    print(input_ids.tolist())
    
    with torch.no_grad():
        embeddings = model(input_ids['input_ids'], attention_mask=input_ids['attention_mask'])[0]

    # calcular a similaridade entre os embeddings
    cos_sim = torch.nn.functional.cosine_similarity(embeddings[0], embeddings[1], dim=0)

    # imprimir o resultado
    st.write("A similaridade entre as frases é:", cos_sim.item())