Views
No views yet
1import torch
2from transformers import AutoTokenizer, XLMRobertaForSequenceClassification
3
4MODEL_PATH = "upb-nlp/xlm_roberta_large_article_same_topic_classification"
5
6tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
7model = XLMRobertaForSequenceClassification.from_pretrained(MODEL_PATH, num_labels=2).to('cuda')
8model.eval()
9
10t1 = "Article title. Article body."
11t2 = "Article title. Article body."
12
13inputs = tokenizer(
14 t1,
15 t2,
16 return_tensors="pt",
17 truncation=True,
18 padding='max_length',
19 max_length=512
20).to('cuda')
21
22# Generate prediction
23with torch.no_grad():
24 outputs = model(**inputs)
25 logits = outputs.logits
26 predicted_class = torch.argmax(logits, dim=1).item()
27
28print(predicted_class)