saved_model_body / load.py
Akhil0-o's picture
Create load.py
9975615
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