import gradio as gr import requests import numpy as np import torch import torch.nn as nn import torchvision.transforms as transforms # from badnet_m import BadNet import timm # model = timm.create_model("hf_hub:nateraw/resnet18-random", pretrained=True) # model.train() # model = BadNet(3, 10) # pipeline = pipeline.to('cuda:0') import os import logging from transformers import WEIGHTS_NAME,AdamW,AlbertConfig,AlbertTokenizer,BertConfig,BertTokenizer from pabee.modeling_albert import AlbertForSequenceClassification from pabee.modeling_bert import BertForSequenceClassification from transformers import glue_output_modes as output_modes from transformers import glue_processors as processors import datasets from whitebox_utils.classifier import MyClassifier import ssl ssl._create_default_https_context = ssl._create_unverified_context import random import numpy as np import torch import argparse def random_seed(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) logger = logging.getLogger(__name__) # TODO: dataset model tokenizer etc. best_model_path = { 'albert_STS-B':'./outputs/train/albert/SST-2/checkpoint-7500', } MODEL_CLASSES = { "bert": (BertConfig, BertForSequenceClassification, BertTokenizer), "albert": (AlbertConfig, AlbertForSequenceClassification, AlbertTokenizer), } model = 'albert' dataset = 'STS-B' task_name = f'{dataset}'.lower() if task_name not in processors: raise ValueError("Task not found: %s" % (task_name)) processor = processors[task_name]() # transformers package-preprocessor output_mode = output_modes[task_name] # output type label_list = processor.get_labels() num_labels = len(label_list) output_dir = f'./PABEE/outputs/train/{model}/{dataset}' data_dir = f'./PABEE/glue_data/{dataset}' config_class, model_class, tokenizer_class = MODEL_CLASSES[model] tokenizer = tokenizer_class.from_pretrained(output_dir, do_lower_case=True) model = model_class.from_pretrained(best_model_path[f'{model}_{dataset}']) exit_type='patience' exit_value=3 classifier = MyClassifier(model,tokenizer,label_list,output_mode,exit_type,exit_value,model) def greet(text,text2,exit_pos): text_input = [(text,text2)] classifier.get_prob_time(text_input,exit_position=exit_pos) iface = gr.Interface(fn=greet, inputs='text', outputs="image") iface.launch()