import csv from dataclasses import dataclass import io import json import logging import random import sys from typing import Dict, List import pandas as pd import streamlit as st import torch import transformers from tqdm import tqdm from autoprompt import utils import autoprompt.create_trigger as ct # logging.getLogger("streamlit.caching").addHandler(logging.StreamHandler(sys.stdout)) # logging.getLogger("streamlit.caching").setLevel(logging.DEBUG) logger = logging.getLogger(__name__) with open('assets/sst2_train.jsonl', 'r') as f: DEFAULT_TRAIN = [json.loads(line) for line in f] @dataclass class CacheTest: """ Stores whether the train button has been pressed for a given set of inputs to run_autoprompt. """ is_test: bool class CacheMiss(Exception): pass def css_hack(): """ Inject some style into this app. ヽ(⌐■_■)ノ """ st.markdown( """ """, unsafe_allow_html=True ) # Setting eq and frozen ensures that a __hash__ method is generated which is needed for caching to # properly respond to changed args. @dataclass(eq=True, frozen=True) class Args: # Configurable template: str model_name: str iters: int num_cand: int accumulation_steps: int # Non-Configurable seed = 0 sentence_size = 64 tokenize_labels = True filter = False initial_trigger = None label_field = "label" bsz = 32 eval_size = 1 @classmethod def from_streamlit(cls): st.sidebar.image('assets/icon.png', width=150) st.sidebar.markdown('### Training Parameters') model_name = st.sidebar.selectbox( "Model", options=['roberta-large', 'bert-base-cased'], help="Language model used for training and evaluation." ) iters = int(st.sidebar.number_input( "Iterations", value=10, min_value=1, max_value=100, help="Number of trigger search iterations. Larger values may yield better results." )) num_cand = int(st.sidebar.number_input( "Number of Candidates", value=25, min_value=1, max_value=100, help="Number of candidate trigger token replacements to evaluate during each search " "iteration. Larger values may yield better results." )) accumulation_steps = int(st.sidebar.number_input( "Gradient Accumulation Steps", value=1, min_value=1, max_value=10, help="Number of gradient accumulation steps used during training. Larger values may yield " "better results. Cannot be larger than half the dataset size." )) st.sidebar.markdown( """ ### Template Templates define how task-specific inputs are combined with trigger tokens to create the prompt. They should contain the following placeholders: - `{sentence}`: Placeholders for the task-specific input fields contain the field name between curly brackets. For manually entered data the field name is `{sentence}`. For uploaded csv's, field names should correspond to columns in the csv. - `[T]`: Placeholder for a trigger token. These are learned from the training data. - `[P]`: Placeholder for where to insert the [MASK] token that the model will predict on. Templates can also include manually written text (such as the period in the default example below). """ ) template = st.sidebar.text_input("Template", "{sentence} [T] [T] [T] [P].") return cls( template=template, model_name=model_name, iters=iters, num_cand=num_cand, accumulation_steps=accumulation_steps, ) # TODO(rloganiv): This probably could use a better name... @dataclass class GlobalData: device: torch.device config: transformers.PretrainedConfig model: transformers.PreTrainedModel tokenizer: transformers.PreTrainedTokenizer embeddings: torch.nn.Module embedding_gradient: ct.GradientStorage predictor: ct.PredictWrapper @classmethod @st.cache(allow_output_mutation=True) def from_pretrained(cls, model_name): logger.info(f'Loading pretrained model: {model_name}') device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') config, model, tokenizer = ct.load_pretrained(model_name) model.to(device) embeddings = ct.get_embeddings(model, config) embedding_gradient = ct.GradientStorage(embeddings) predictor = ct.PredictWrapper(model) return cls( device, config, model, tokenizer, embeddings, embedding_gradient, predictor ) @dataclass class Dataset: train: List[int] label_map: Dict[str, str] def load_trigger_dataset(dataset, templatizer): instances = [] for x in dataset: instances.append(templatizer(x)) return instances @st.cache(suppress_st_warning=True, allow_output_mutation=True, hash_funcs={CacheTest: lambda o: 0}) def run_autoprompt(args, dataset, cache_test): if cache_test.is_test: raise CacheMiss() ct.set_seed(args.seed) global_data = GlobalData.from_pretrained(args.model_name) templatizer = utils.TriggerTemplatizer( args.template, global_data.config, global_data.tokenizer, label_field=args.label_field, label_map=dataset.label_map, tokenize_labels=args.tokenize_labels, add_special_tokens=True, ) evaluation_fn = ct.AccuracyFn(global_data.tokenizer, dataset.label_map, global_data.device, tokenize_labels=args.tokenize_labels) # Do not allow for initial trigger specification. trigger_ids = [global_data.tokenizer.mask_token_id] * templatizer.num_trigger_tokens trigger_ids = torch.tensor(trigger_ids, device=global_data.device).unsqueeze(0) best_trigger_ids = trigger_ids.clone() # Load datasets logger.info('Loading datasets') collator = utils.Collator(pad_token_id=global_data.tokenizer.pad_token_id) try: train_dataset = load_trigger_dataset(dataset.train, templatizer) except KeyError as e: raise RuntimeError( 'A field in your template is not present in the uploaded dataset. ' f'Check that there is a column with the name: {e}' ) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=args.bsz, shuffle=True, collate_fn=collator) progress = st.progress(0.0) trigger_placeholder = st.empty() best_dev_metric = -float('inf') for i in range(args.iters): logger.info(f'Iteration: {i}') progress.progress(float(i)/args.iters) current_trigger = ','.join(global_data.tokenizer.convert_ids_to_tokens(best_trigger_ids.squeeze(0))) trigger_placeholder.markdown(f'**Current trigger**: {current_trigger}') global_data.model.zero_grad() train_iter = iter(train_loader) averaged_grad = None # Compute gradient of loss for step in range(args.accumulation_steps): try: model_inputs, labels = next(train_iter) except: logger.warning( 'Insufficient data for number of accumulation steps. ' 'Effective batch size will be smaller than specified.' ) break model_inputs = {k: v.to(global_data.device) for k, v in model_inputs.items()} labels = labels.to(global_data.device) predict_logits = global_data.predictor(model_inputs, trigger_ids) loss = ct.get_loss(predict_logits, labels).mean() loss.backward() grad = global_data.embedding_gradient.get() bsz, _, emb_dim = grad.size() selection_mask = model_inputs['trigger_mask'].unsqueeze(-1) grad = torch.masked_select(grad, selection_mask) grad = grad.view(bsz, templatizer.num_trigger_tokens, emb_dim) if averaged_grad is None: averaged_grad = grad.sum(dim=0) / args.accumulation_steps else: averaged_grad += grad.sum(dim=0) / args.accumulation_steps logger.info('Evaluating Candidates') pbar = tqdm(range(args.accumulation_steps)) train_iter = iter(train_loader) token_to_flip = i % templatizer.num_trigger_tokens candidates = ct.hotflip_attack(averaged_grad[token_to_flip], global_data.embeddings.weight, increase_loss=False, num_candidates=args.num_cand) current_score = 0 candidate_scores = torch.zeros(args.num_cand, device=global_data.device) denom = 0 for step in pbar: try: model_inputs, labels = next(train_iter) except: logger.warning( 'Insufficient data for number of accumulation steps. ' 'Effective batch size will be smaller than specified.' ) break model_inputs = {k: v.to(global_data.device) for k, v in model_inputs.items()} labels = labels.to(global_data.device) with torch.no_grad(): predict_logits = global_data.predictor(model_inputs, trigger_ids) eval_metric = evaluation_fn(predict_logits, labels) # Update current score current_score += eval_metric.sum() denom += labels.size(0) # NOTE: Instead of iterating over tokens to flip we randomly change just one each # time so the gradients don't get stale. for i, candidate in enumerate(candidates): # if candidate.item() in filter_candidates: # candidate_scores[i] = -1e32 # continue temp_trigger = trigger_ids.clone() temp_trigger[:, token_to_flip] = candidate with torch.no_grad(): predict_logits = global_data.predictor(model_inputs, temp_trigger) eval_metric = evaluation_fn(predict_logits, labels) candidate_scores[i] += eval_metric.sum() if (candidate_scores >= current_score).any(): logger.info('Better trigger detected.') best_candidate_score = candidate_scores.max() best_candidate_idx = candidate_scores.argmax() trigger_ids[:, token_to_flip] = candidates[best_candidate_idx] logger.info(f'Train metric: {best_candidate_score / (denom + 1e-13): 0.4f}') # Skip eval best_trigger_ids = trigger_ids.clone() progress.progress(1.0) current_trigger = ','.join(global_data.tokenizer.convert_ids_to_tokens(best_trigger_ids.squeeze(0))) trigger_placeholder.markdown(f'**Current trigger**: {current_trigger}') best_trigger_tokens = global_data.tokenizer.convert_ids_to_tokens(best_trigger_ids.squeeze(0)) train_output = predict_test(map(lambda x: x['sentence'], dataset.train), dataset.label_map, templatizer, best_trigger_ids, global_data.tokenizer, global_data.predictor, args) # Streamlit does not like accessing widgets across functions, which is # problematic for this "live updating" widget which we want to still # display even if the train output is cached. To get around this, we're # going to delete the widget and replace it with a very similar looking # widget outside the function...no one will ever notice ;) trigger_placeholder.empty() return ( best_trigger_tokens, current_score/denom, dataset.label_map, templatizer, best_trigger_ids, global_data.tokenizer, global_data.predictor, args, train_output ) def predict_test(sentences, label_map, templatizer, best_trigger_ids, tokenizer, predictor, args): # Evaluate clean output = { 'sentences': [] } any_label = None for label in label_map.values(): output[label] = [] any_label = label output['prompt'] = [] for sentence in sentences: model_inputs, _ = templatizer({'sentence': sentence, 'label': any_label}) model_inputs = {k: v.to(best_trigger_ids.device) for k, v in model_inputs.items()} prompt_ids = ct.replace_trigger_tokens( model_inputs, best_trigger_ids, model_inputs['trigger_mask']) prompt = ' '.join(tokenizer.convert_ids_to_tokens(prompt_ids['input_ids'][0])) output['prompt'].append(prompt) predict_logits = predictor(model_inputs, best_trigger_ids) output['sentences'].append(sentence) for label in label_map.values(): label_id = utils.encode_label(tokenizer=tokenizer, label=label, tokenize=args.tokenize_labels) label_id = label_id.to(best_trigger_ids.device) label_loss = ct.get_loss(predict_logits, label_id) # st.write(sentence, label, label_loss) output[label].append(label_loss.item()) return output def manual_dataset(use_defaults): num_train_instances = st.slider("Number of Train Instances", 4, 32, 8) any_empty = False dataset = [] data_col, label_col = st.beta_columns([3,1]) for i in range(num_train_instances): default_data = DEFAULT_TRAIN[i]['sentence'] if use_defaults else '' default_label = DEFAULT_TRAIN[i]['label'] if use_defaults else '' with data_col: data = st.text_input("Train Instance " + str(i+1), default_data) with label_col: label = st.text_input("Train Label " + str(i+1), default_label, max_chars=20) if data == "" or label == "": any_empty = True dataset.append({'sentence': data, 'label': label}) label_set = list(set(map(lambda x: x['label'], dataset))) label_idx = {x: i for i, x in enumerate(label_set)} label_map = dict(map(lambda x: (x, x), label_set)) if any_empty: st.warning('Waiting for data to be added') st.stop() if len(label_set) < 2: st.warning('Not enough labels') st.stop() return Dataset( train=dataset, label_map=label_map ) def csv_dataset(): st.markdown(""" Please upload your training and evaluation csv files. Format restrictions: - The file is required to have a header - The column name of the output field should be `label`. - Each file should contain no more than 64 rows. """) train_csv = st.file_uploader('Train', accept_multiple_files=False) if train_csv is None: st.stop() with io.StringIO(train_csv.getvalue().decode('utf-8')) as f: reader = csv.DictReader(f) train_dataset = list(reader) if len(train_dataset) > 64: raise ValueError('Train dataset is too large. Please limit the number ' 'of examples to 64 or less.') labels = set(x['label'] for x in train_dataset) label_map = {x: x for x in labels} return Dataset( train=train_dataset, label_map=label_map ) def run(): css_hack() st.title('AutoPrompt Demo') st.markdown(''' For many years, the predominant approach for training machine learning models to solve NLP tasks has been to use supervised training data to estimate model parameters using maximum likelihood estimation or some similar paradigm. Whether fitting a logistic regression model over a bag-of-words, an LSTM over a sequence of GloVe embeddings, or finetuning a language model such as ELMo or BERT, the approach is essentially the same. However, as language models have become more and more capable of accurately generating plausible text a new possibility for solving classification tasks has emerged... ## Prompting Prompting is the method of converting classification tasks into *fill-in-the-blanks* problems that can be solved by a language model **without modifying the model's internals**. For example, to perform sentiment analysis, we may take the sentence we wish to classify and append the text "Overall, this movie was ____." and feed it into a language model like so: ''') st.image('assets/bert-mouth.png', use_column_width=True) st.markdown(''' By measuring whether the language model assigns a higher probability to words that are associated with a **positive** sentiment ("good", "great", and "fantastic") vs. words that are associated with a **negative** sentiment ("bad", "terrible", or "awful") we can infer the predicted label for the given input. So in this example, because the word "good" has a higher probability than "bad", the predicted label is **positive**. ## AutoPrompt One issue that arises when using prompts is that it is not usually clear how to best pose a task as a fill-in-the-blanks problem in a way that gets the most performance from the language model. Even for a simple problem like sentiment analysis, we don't know whether it is better to ask whether a movie is good/bad, or whether you feel great/terrible about it, and for more abstract problems like natural language inference it is difficult to even know where to start. To cure this writer's block we introduce **AutoPrompt**, a data-driven approach for automatic prompt construction. The basic idea is straightfoward: instead of writing a prompt, a user need only write a **template** that specfies where the *task inputs* go along with placeholders for a number of *trigger tokens* that will automatically be learned by the model and the *predict token* that the model will fill in: ''') st.image('assets/template.png', use_column_width=True) st.markdown( ''' In each iteration of the search process: 1. The template is instantiated using a batch of training inputs. 2. The loss of the model on each input is measured and used to identify a number of candidate replacements for the current trigger tokens. 3. The performance of each candidate is measured on another batch of training data, and the best performing candidate is used in the next iteration. ### Demo To give a better sense of how AutoPrompt works, we have provided a simple interactive demo. You can generate a prompt using the training data we have pre-populated for you, or alternatively write your own training/evaluation instances or upload them using a csv below. In addition, you can vary some of the training parameters, as well as the template using the sidebar on the left. ''' ) args = Args.from_streamlit() dataset_mode = st.radio('How would you like to input your training data?', options=['Example Data', 'Manual Input', 'From CSV']) if dataset_mode == 'Example Data': dataset = manual_dataset(use_defaults=True) elif dataset_mode == 'Manual Input': dataset = manual_dataset(use_defaults=False) else: dataset = csv_dataset() button = st.empty() clicked = button.button('Train') if clicked: trigger_tokens, eval_metric, label_map, templatizer, best_trigger_ids, tokenizer, predictor, args, train_output = run_autoprompt(args, dataset, cache_test=CacheTest(False)) else: try: trigger_tokens, eval_metric, label_map, templatizer, best_trigger_ids, tokenizer, predictor, args, train_output = run_autoprompt(args, dataset, cache_test=CacheTest(True)) except CacheMiss: st.stop() else: button.empty() st.markdown(f'**Final trigger**: {", ".join(trigger_tokens)}') st.dataframe(pd.DataFrame(train_output).style.highlight_min(axis=1, color='#94666b')) logger.debug('Dev metric') st.write('Accuracy: ' + str(round(eval_metric.item()*100, 1))) st.write(""" Et voila, you've now effectively finetuned a classifier using just a few kilobytes of parameters (the tokens in the prompt). If you like you can write down your "model" on the back of a napkin and take it with you. ### Try it out yourself! """) sentence = st.text_input("Sentence", 'Enter a test input here') pred_output = predict_test([sentence], label_map ,templatizer, best_trigger_ids, tokenizer, predictor, args) st.dataframe(pd.DataFrame(pred_output).style.highlight_min(axis=1, color='#94666b')) st.markdown(''' ## Where can I learn more? If you are interested in learning more about AutoPrompt we recommend [reading our paper](https://arxiv.org/abs/2010.15980) and [checking out our code](https://github.com/ucinlp/autoprompt), or if you'd like you can also watch our presentation at EMNLP 2020: ''') st.components.v1.iframe( src="https://www.youtube.com/embed/IBMT_oOCBbc", height=400, ) st.markdown('Thanks!') if __name__ == '__main__': logging.basicConfig(level=logging.INFO, stream=sys.stdout) run()