import streamlit as st
import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = 'intfloat/multilingual-e5-large'
adapters_name = './checkpoint-21170'


model = AutoModelForCausalLM.from_pretrained(model_name)
model = PeftModel.from_pretrained(model, adapters_name)
model = model.merge_and_unload()

tokenizer = AutoTokenizer.from_pretrained(model_name)

description = st.text_input("Product description")
review = st.text_input("Review")

if description and review:
    input_texts = [
        f'query: {review}',
        f'passage: {description}'
    ]
    batch_dict = tokenizer(input_texts, max_length=512,
                           padding=True, truncation=True, return_tensors='pt')

    query_embedding, doc_embedding = model(**batch_dict).pooler_output

    similarity = torch.nn.functional.cosine_similarity(
        query_embedding, doc_embedding)

    threshold = 0.7

    if similarity > threshold:
        st.write('Relevant')
    else:
        st.write('Irrelevant')