Views
No views yet
1import torch
2from transformers import AutoModelForSequenceClassification, AutoTokenizer
3
4# label_list
5label_list = ['matched', 'unmatched']
6
7# Load model from HuggingFace Hub
8tokenizer = AutoTokenizer.from_pretrained("Fan-s/reddit-tc-bert", use_fast=True)
9model = AutoModelForSequenceClassification.from_pretrained("Fan-s/reddit-tc-bert")
10
11# Set the input
12post = "don't make gravy with asbestos."
13response = "i'd expect someone with a culinary background to know that. since we're talking about school dinner ladies, they need to learn this pronto."
14
15# Predict whether the two sentences are matched
16def predict(post, response, max_seq_length=128):
17 with torch.no_grad():
18 args = (post, response)
19 input = tokenizer(*args, padding="max_length", max_length=max_seq_length, truncation=True, return_tensors="pt")
20 output = model(**input)
21 logits = output.logits
22 item = torch.argmax(logits, dim=1)
23 predict_label = label_list[item]
24 return predict_label, logits
25
26predict_label, logits = predict(post, response)
27# Matched
28print("predict_label:", predict_label)