Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForSequenceClassification
3
4classifier_model_path = "jmvcoelho/ad-classifier-v0.0"
5tokenizer = AutoTokenizer.from_pretrained(classifier_model_path)
6model = AutoModelForSequenceClassification.from_pretrained(classifier_model_path)
7model.eval()
8
9device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
10model.to(device)
11
12def classify(passages):
13 inputs = tokenizer(
14 passages, padding=True, truncation=True, max_length=512, return_tensors="pt"
15 )
16 inputs = {k: v.to(device) for k, v in inputs.items()}
17 with torch.no_grad():
18 outputs = model(**inputs)
19 logits = outputs.logits
20 predictions = torch.argmax(logits, dim=-1)
21 return predictions.cpu().tolist()
22
23preds = classify(["sample_text_1", "sample_text_2"])