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