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