Views
No views yet
1import torch.nn.functional as F
2from transformers import AutoTokenizer, AutoModelForSequenceClassification
3tokenizer = AutoTokenizer.from_pretrained("assemblyai/distilbert-base-uncased-qqp")
4model = AutoModelForSequenceClassification.from_pretrained("assemblyai/distilbert-base-uncased-qqp")
5
6tokenized_segments = tokenizer(["How many hours does it take to fly from California to New York?"], ["What is the flight time from New York to Seattle?"], return_tensors="pt", padding=True, truncation=True)
7tokenized_segments_input_ids, tokenized_segments_attention_mask = tokenized_segments.input_ids, tokenized_segments.attention_mask
8model_predictions = F.softmax(model(input_ids=tokenized_segments_input_ids, attention_mask=tokenized_segments_attention_mask)['logits'], dim=1)
9
10print("Duplicate probability: "+str(model_predictions[0][1].item()*100)+"%")
11print("Non-duplicate probability: "+str(model_predictions[0][0].item()*100)+"%")