Spaces:
Running
Running
| from nltk.tokenize import sent_tokenize | |
| import pandas as pd | |
| ###################### | |
| # prerequisite: | |
| # 1. Pip install transformer | |
| # 2. Define tokenizer + MAX_LEN | |
| # 3. Construct DistillBERTClass_SL class | |
| # 4. Construct Triage_SL class | |
| # 5. Define predict__SL class | |
| # 6. Load model_SL & call eval() | |
| # 7. Pre_define predict_params_SL | |
| #################### | |
| from transformers import DistilBertTokenizer | |
| tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-cased') | |
| import torch | |
| """### DataSet Class -- Triage_SL""" | |
| from torch.utils.data import Dataset, DataLoader | |
| class Triage_SL(Dataset): | |
| # initialize the directory containing the dataframe, the tokenizer, and the max lens of sentences | |
| def __init__(self, dataframe, tokenizer, max_len): | |
| self.len = len(dataframe) | |
| self.data = dataframe | |
| self.tokenizer = tokenizer # load in tokenizer, used in _getitem | |
| self.max_len = max_len | |
| # The __getitem__ function loads and returns a sample from the dataset at the given index idx. | |
| def __getitem__(self, index): | |
| if index >= len(self): | |
| raise StopIteration | |
| # preprossessing sentences to standarize format as in: word+""+word | |
| sent = str(self.data.sentence[index]) | |
| sent = " ".join(sent.split()) | |
| # 1.- Split the sentence into tokens. | |
| # 2.- Add the special [CLS] and [SEP] tokens. | |
| # 3.- Map the tokens to their IDs. | |
| # 4.- Pad or truncate all sentences to the same length. | |
| # 5.- Create the attention masks which explicitly differentiate real tokens from [PAD] tokens. | |
| inputs = self.tokenizer.encode_plus( | |
| sent, # Sentence to encode | |
| None, # text_pair | |
| add_special_tokens=True, # Add '[CLS]' and '[SEP]' | |
| max_length=self.max_len, | |
| pad_to_max_length=True, # Pad & truncate all sentences. | |
| return_token_type_ids=True, | |
| truncation=True | |
| ) | |
| ids = inputs['input_ids'] | |
| mask = inputs['attention_mask'] | |
| return { | |
| 'ids': torch.tensor(ids, dtype=torch.long), | |
| 'mask': torch.tensor(mask, dtype=torch.long), | |
| # 'targets': torch.tensor(self.data.ENCODE_LABEL[index], dtype=torch.float), # sentence label -> y value | |
| # 'combined_label': self.data.combined_label[index] | |
| } | |
| # The __len__ function returns the number of samples in our dataset. | |
| def __len__(self): | |
| return self.len | |
| # read in an essay and resturns a df in sentence level | |
| def essay_to_sent_df(essay): | |
| sentences = [] | |
| paragraphs = [l for l in essay.split('\n') if len(l) > 0] | |
| for para in paragraphs: | |
| # tokenize paragraph by "." and concatenate to sentences[] | |
| sentences.extend(sent_tokenize(para)) | |
| return pd.DataFrame(sentences, columns=['sentence']) | |
| # Defining some key variables that will be used later on in the training | |
| MAX_LEN = 512 | |
| """### Predefine predict_params_SL""" | |
| PREDICT_BATCH_SIZE = 1 | |
| predict_params_SL = {'batch_size': PREDICT_BATCH_SIZE, | |
| 'shuffle': False, | |
| 'num_workers': 0 | |
| } | |
| """### Predict Fn -- predict_SL""" | |
| sigmoid = torch.nn.Sigmoid() | |
| def predict_SL(model, validation_loader): | |
| epoch_val_outputs=[] | |
| cpu_device = 'cpu' | |
| model.eval() | |
| with torch.no_grad(): | |
| for _, data in enumerate(validation_loader, 0): | |
| ids = data['ids'].to(cpu_device, dtype = torch.long) | |
| mask = data['mask'].to(cpu_device, dtype = torch.long) | |
| outputs = model(ids, mask)["logits"].squeeze() # ??squeeze?? | |
| outputs = (sigmoid(outputs).data>0.5).float() | |
| epoch_val_outputs.append(outputs.item()) | |
| return epoch_val_outputs | |
| def predict_mainidea_sent_old(paragraph, model): | |
| # prepare data | |
| sent_df = essay_to_sent_df(paragraph) | |
| predicting_SL_set = Triage_SL(sent_df, tokenizer, MAX_LEN) | |
| predicting_SL_loader = DataLoader(predicting_SL_set, **predict_params_SL) | |
| # load model to device | |
| device = 'cpu' | |
| model.to(device) | |
| # predict + roundup | |
| sent_label = predict_SL(model, predicting_SL_loader) | |
| print(sent_label) | |
| return pd.DataFrame([(str(l), s) for l, s in zip(sent_label, sent_df.sentence)], columns=['label', 'sentence']) | |