Spaces:
Runtime error
Runtime error
| import numpy as np | |
| import gradio as gr | |
| from transformers import AutoTokenizer,AutoModelForSequenceClassification | |
| from transformers import set_seed | |
| from torch.utils.data import Dataset,DataLoader | |
| import torch | |
| import torch.nn as nn | |
| import numpy as np | |
| import warnings | |
| warnings.filterwarnings('ignore') | |
| set_seed(4) | |
| device = "cpu" | |
| model_checkpoint = "facebook/esm2_t6_8M_UR50D" | |
| tokenizer = AutoTokenizer.from_pretrained(model_checkpoint) | |
| def AMP(file): | |
| test_sequences = file | |
| class MyDataset(Dataset): | |
| def __init__(self,dict_data) -> None: | |
| super(MyDataset,self).__init__() | |
| self.data=dict_data | |
| def __getitem__(self, index): | |
| return self.data['text'][index] | |
| def __len__(self): | |
| return len(self.data['text']) | |
| def collate_fn(batch): | |
| max_len = 30 | |
| pt_batch=tokenizer([b[0] for b in batch], max_length=max_len, padding="max_length",truncation=True, return_tensors='pt') | |
| return {'input_ids':pt_batch['input_ids'], | |
| 'attention_mask':pt_batch['attention_mask']} | |
| test_dict = {"text":test_sequences} | |
| test_data=MyDataset(test_dict) | |
| test_dataloader=DataLoader(test_data,batch_size=len(test_sequences), collate_fn=collate_fn) | |
| class MyModel(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.bert = AutoModelForSequenceClassification.from_pretrained(model_checkpoint,num_labels=320) | |
| self.bn1 = nn.BatchNorm1d(256) | |
| self.bn2 = nn.BatchNorm1d(128) | |
| self.bn3 = nn.BatchNorm1d(64) | |
| self.relu = nn.ReLU() | |
| self.fc1 = nn.Linear(320,256) | |
| self.fc2 = nn.Linear(256,128) | |
| self.fc3 = nn.Linear(128,64) | |
| self.output_layer = nn.Linear(64,2) | |
| self.dropout = nn.Dropout(0) | |
| def forward(self,x): | |
| with torch.no_grad(): | |
| bert_output = self.bert(input_ids=x['input_ids'].to(device),attention_mask=x['attention_mask'].to(device)) | |
| output_feature = self.dropout(bert_output["logits"]) | |
| output_feature = self.relu(self.bn1(self.fc1(output_feature))) | |
| output_feature = self.relu(self.bn2(self.fc2(output_feature))) | |
| output_feature = self.relu(self.bn3(self.fc3(output_feature))) | |
| output_feature = self.output_layer(output_feature) | |
| return torch.softmax(output_feature,dim=1) | |
| model = MyModel() | |
| model.load_state_dict(torch.load("best_model.pth")) | |
| model = model.to(device) | |
| model.eval() | |
| out = [] | |
| out_probability = [] | |
| with torch.no_grad(): | |
| for index, batch in enumerate(test_dataloader): | |
| batchs = {k: v for k, v in batch.items()} | |
| predict = model(batchs) | |
| # tsne_plot(output_feature,batchs) | |
| out_probability.extend(np.max(np.array(predict.cpu()),axis=1).tolist()) | |
| # out_probability.extend(np.array(predict.cpu())[:, -1].tolist()) | |
| test_argmax = np.argmax(predict.cpu(), axis=1).tolist() | |
| out.extend(test_argmax) | |
| return out, out_probability | |
| iface = gr.Interface(fn=AMP, | |
| inputs="text"(label="input sequence"), | |
| outputs= ["text"(label="class"), "text"(label="probability")]) | |
| iface.launch() |