Views
No views yet
1from transformers import BertTokenizer, BertForSequenceClassification
2import torch
3
4tokenizer = BertTokenizer.from_pretrained('fzn0x/bert-spam-classification-model')
5model = BertForSequenceClassification.from_pretrained('fzn0x/bert-spam-classification-model')
6
7device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
8model.to(device)
9model.eval()
10
11def model_predict(text: str):
12 inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)
13 with torch.no_grad():
14 outputs = model(**inputs)
15 logits = outputs.logits
16 prediction = torch.argmax(logits, dim=1).item()
17 return 'SPAM' if prediction == 1 else 'HAM'
18
19def predict():
20 text = "Hello, do you know with this crypto you can be rich? contact us in 88888"
21 predicted_label = model_predict(text)
22 print(f"1. Predicted class: {predicted_label}") # EXPECT: SPAM
23
24 text = "Help me richard!"
25 predicted_label = model_predict(text)
26 print(f"2. Predicted class: {predicted_label}") # EXPECT: HAM
27
28 text = "You can buy loopstation for 100$, try buyloopstation.com"
29 predicted_label = model_predict(text)
30 print(f"3. Predicted class: {predicted_label}") # EXPECT: SPAM
31
32 text = "Mate, I try to contact your phone, where are you?"
33 predicted_label = model_predict(text)
34 print(f"4. Predicted class: {predicted_label}") # EXPECT: HAM
35
36if __name__ == "__main__":
37 predict()citations.bib for full BibTeX entries.