utils

seal/run_inference.py

@@ -0,0 +1,82 @@
+from unittest import result
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, Subset
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from datasets import load_dataset
+
+import os
+import numpy as np
+from tqdm import tqdm
+
+from utils.inference_utils import InferenceResults, saveResults
+
+# Load validation set
+
+def load_session(dataset, model, split):
+    dataset = load_dataset(dataset, split=split)
+    dataloader = DataLoader(
+        dataset,
+        batch_size=256, drop_last=True
+    )
+    model = AutoModelForSequenceClassification.from_pretrained(model)
+    tokenizer = AutoTokenizer.from_pretrained(model)
+    return tokenizer, dataloader, model
+
+# Add hook to capture hidden layer
+def get_input(name, model):
+    hidden_layers = {}
+    def hook(model, input, output):
+        if name in hidden_layers:
+            del hidden_layers[name]
+        hidden_layers[name] = input[0].detach()
+    return hook, hidden_layers
+
+def run_inference(dataset='yelp_polarity', model='textattack/albert-base-v2-yelp-polarity', split='test', output_path='./assets/data/inference_results'):
+    tokenizer, dataloader, model = load_session(dataset,model,split)
+    model.eval()
+    model.to('cpu')
+    hook, hidden_layers = model.classifier.register_forward_hook(get_input('last_layer', model))
+    # Run inference on entire dataset
+    hidden_list = []
+    loss_list = []
+    output_list = []
+    example = []
+    labels = []
+    criterion = nn.CrossEntropyLoss(reduction='none')
+    softmax = nn.Softmax(dim=1)
+    with torch.no_grad():
+        for batch_num, batch in tqdm(enumerate(dataloader), total=len(dataloader), position=0, leave=True):
+            batch_ex = [ex[:512] for ex in batch['text']]
+            inputs = tokenizer(batch_ex, padding=True, return_tensors='pt').to('cpu')
+            targets = batch['label']
+            
+            outputs = model(**inputs)['logits']
+            loss = criterion(outputs, targets)
+            predictions = softmax(outputs)
+            
+            hidden_list.append(hidden_layers['last_layer'].cpu())
+            loss_list.append(loss.cpu())
+            #output_list.append(predictions[:, 1].cpu())
+            output_list.append(np.argmax(predictions, axis=1))
+            labels.append(targets)
+            example.append(inputs['input_ids'])
+    embeddings = torch.vstack(hidden_list)
+    #outputs = torch.hstack(output_list)
+    losses = torch.hstack(loss_list)
+    targets = torch.hstack(labels)
+    #inputs = torch.hstack(example)
+    results = save_results(embeddings,losses,targets)
+    saveResults(os.path.join(output_path,dataset+'.pkl'),results)
+
+
+
+def save_results(embeddings, losses, labels):
+    results = InferenceResults(
+        embeddings = torch.clone(embeddings),
+        losses     = losses,
+        labels     = labels
+    )
+    return results
+    
+

seal/utils/__init__.py

@@ -0,0 +1 @@
+from .style_hacks import *

seal/utils/inference_utils.py

@@ -0,0 +1,24 @@
+import pickle
+from dataclasses import dataclass
+import torch
+
+@dataclass
+class InferenceResults:
+    """
+    Class for storing embeddings and losses from running inference on a model.
+    
+    Fields:
+    - embeddings: (num_examples x num_dimensions) tensor of last-layer embeddings
+    - losses: (num_examples x 1) tensor of losses
+    - outputs: optional (num_examples x num_classes) tensor of output logits
+    - labels: optional (num_examples x 1) tensor of labels
+    """
+    
+    embeddings: torch.Tensor
+    losses: torch.Tensor
+    outputs: torch.Tensor = None
+    labels: torch.Tensor = None 
+
+def saveResults(fname, results):
+    with open(fname, 'wb+') as f:
+        pickle.dump(results, f)

seal/utils/style_hacks.py

@@ -0,0 +1,86 @@
+"""
+ placeholder for all streamlit style hacks
+"""
+import streamlit as st
+
+
+def init_style():
+    return st.markdown(
+        """
+    <style>
+    /* Side Bar */
+    [data-testid="stSidebar"][aria-expanded="true"] > div:first-child {
+         width: 250px;
+       }
+    [data-testid="stSidebar"][aria-expanded="false"] > div:first-child {
+        width: 250px;
+    }
+    [data-testid="stSidebar"]{
+        flex-basis: unset;
+    }
+    .css-1outpf7 {
+        background-color:rgb(254 244 219);
+        width:10rem;
+        padding:10px 10px 10px 10px;
+    }
+
+    /* Main Panel*/  
+    .css-18e3th9 {
+        padding:10px 10px 10px -200px;
+    }
+    .css-1ubw6au:last-child{
+        background-color:lightblue;
+    }
+
+    /* Model Panels : element-container */
+    .element-container{
+            border-style:none
+    }
+
+    /* Radio Button Direction*/
+    div.row-widget.stRadio > div{flex-direction:row;}
+
+    /* Expander Boz*/
+    .streamlit-expander {
+        border-width: 0px;
+        border-bottom: 1px solid #A29C9B;
+        border-radius: 10px;
+    }
+
+    .streamlit-expanderHeader {
+        font-style: italic;
+        font-weight :600;
+        font-size:16px;
+        padding-top:0px;
+        padding-left: 0px;
+        color:#A29C9B
+
+    /* Section Headers */
+    .sectionHeader {
+        font-size:10px;
+    }
+    [data-testid="stMarkdownContainer]{
+        font-family: sans-serif;
+        font-weight: 500;
+        font-size: 1.5 rem !important;
+        color: rgb(250, 250, 250);
+        padding: 1.25rem 0px 1rem;
+        margin: 0px;
+        line-height: 1.4;
+    }
+
+    /* text input*/
+    .st-e5 {
+        background-color:lightblue;
+    }
+    /*line special*/
+    .line-one{
+        border-width: 0px;
+        border-bottom: 1px solid #A29C9B;
+        border-radius: 50px; 
+    }
+
+    </style>
+""",
+        unsafe_allow_html=True,
+    )