from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch
import torch.nn.functional as F


MODEL_ID = "prhegde/query-product-relevance-model-ecommerce"
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)

query_text = "sofa with ottoman"
prod_text = "daryl 82 '' wide reversible sofa & chaise with ottoman"

tok_output = tokenizer(query_text, prod_text, padding='max_length',
                      max_length=160, truncation=True, return_tensors='pt',
                      return_attention_mask=True)

input_ids = tok_output.input_ids
attention_masks = tok_output.attention_mask
token_type_ids = tok_output.token_type_ids

output = model(input_ids, attention_mask = attention_masks,
                  token_type_ids = token_type_ids)
probs = F.softmax(output.logits, dim=-1)
score = probs[0][1].item()
print(score)