Spaces:
Running
Running
| from transformers import RobertaTokenizer, RobertaForSequenceClassification, RobertaModel | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from torch.utils.data import TensorDataset, DataLoader | |
| device = torch.device("cpu") | |
| class MLP(nn.Module): | |
| def __init__(self, input_dim): | |
| super(MLP, self).__init__() | |
| self.fc1 = nn.Linear(input_dim, 256) | |
| self.fc2 = nn.Linear(256, 2) | |
| self.gelu = nn.GELU() | |
| def forward(self, x): | |
| x = self.gelu(self.fc1(x)) | |
| x = self.fc2(x) | |
| return x | |
| def extract_features(text): | |
| tokenizer = RobertaTokenizer.from_pretrained("roberta-base") | |
| model = RobertaModel.from_pretrained("roberta-base").to(device) | |
| tokenized_text = tokenizer.encode(text, truncation=True, max_length=512, return_tensors="pt") | |
| outputs = model(tokenized_text) | |
| last_hidden_states = outputs.last_hidden_state | |
| TClassification = last_hidden_states[:, 0, :].squeeze().detach().numpy() | |
| return TClassification | |
| def RobertaSentinelOpenGPTInference(input_text): | |
| features = extract_features(input_text) | |
| loaded_model = MLP(768).to(device) | |
| loaded_model.load_state_dict(torch.load("SentinelCheckpoint/RobertaSentinelOpenGPT.pth", map_location=device)) | |
| # Define the tokenizer and model for feature extraction | |
| with torch.no_grad(): | |
| inputs = torch.tensor(features).to(device) | |
| outputs = loaded_model(inputs.float()) | |
| _, predicted = torch.max(outputs, 0) | |
| Probs = (F.softmax(outputs, dim=0).cpu().numpy()) | |
| return Probs | |
| def RobertaSentinelCSAbstractInference(input_text): | |
| features = extract_features(input_text) | |
| loaded_model = MLP(768).to(device) | |
| loaded_model.load_state_dict(torch.load("SentinelCheckpoint/RobertaSentinelCSAbstract.pth", map_location=device)) | |
| # Define the tokenizer and model for feature extraction | |
| with torch.no_grad(): | |
| inputs = torch.tensor(features).to(device) | |
| outputs = loaded_model(inputs.float()) | |
| _, predicted = torch.max(outputs, 0) | |
| Probs = (F.softmax(outputs, dim=0).cpu().numpy()) | |
| return Probs | |
| def RobertaClassifierOpenGPTInference(input_text): | |
| tokenizer = RobertaTokenizer.from_pretrained("roberta-base") | |
| model_path = "ClassifierCheckpoint/RobertaClassifierOpenGPT.pth" | |
| model = RobertaForSequenceClassification.from_pretrained('roberta-base', num_labels=2) | |
| model.load_state_dict(torch.load(model_path), map_location=device) | |
| model = model.to(torch.device('cpu')) | |
| model.eval() | |
| tokenized_input = tokenizer(input_text, truncation=True, padding=True, max_length=512, return_tensors='pt') | |
| input_ids = tokenized_input['input_ids'].to(torch.device('cpu')) | |
| attention_mask = tokenized_input['attention_mask'].to(torch.device('cpu')) | |
| # Make a prediction | |
| with torch.no_grad(): | |
| outputs = model(input_ids, attention_mask=attention_mask) | |
| logits = outputs.logits | |
| Probs = F.softmax(logits, dim=1).cpu().numpy()[0] | |
| return Probs | |
| def RobertaClassifierCSAbstractInference(input_text): | |
| tokenizer = RobertaTokenizer.from_pretrained("roberta-base") | |
| model_path = "ClassifierCheckpoint/RobertaClassifierCSAbstract.pth" | |
| model = RobertaForSequenceClassification.from_pretrained('roberta-base', num_labels=2) | |
| model.load_state_dict(torch.load(model_path, map_location=device)) | |
| model = model.to(torch.device('cpu')) | |
| model.eval() | |
| tokenized_input = tokenizer(input_text, truncation=True, padding=True, max_length=512, return_tensors='pt') | |
| input_ids = tokenized_input['input_ids'].to(torch.device('cpu')) | |
| attention_mask = tokenized_input['attention_mask'].to(torch.device('cpu')) | |
| # Make a prediction | |
| with torch.no_grad(): | |
| outputs = model(input_ids, attention_mask=attention_mask) | |
| logits = outputs.logits | |
| Probs = F.softmax(logits, dim=1).cpu().numpy()[0] | |
| return Probs | |