|
import torch |
|
from torch.utils.data import Dataset |
|
import pandas as pd |
|
from sklearn.model_selection import train_test_split |
|
from sklearn.metrics import classification_report |
|
from transformers import RobertaTokenizer, RobertaForSequenceClassification, Trainer, TrainingArguments |
|
from transformers import TrainerCallback |
|
import os |
|
from transformers import TrainingArguments, Trainer |
|
|
|
model = RobertaForSequenceClassification.from_pretrained("./best_model") |
|
tokenizer = RobertaTokenizer.from_pretrained("./best_model") |
|
|
|
|
|
def maliciousornot(link): |
|
inputs = tokenizer(link, return_tensors="pt") |
|
outputs = model(**inputs) |
|
predictions = torch.argmax(outputs.logits, dim=-1) |
|
return predictions |
|
|