File size: 5,489 Bytes
b208d2e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 |
from typing import Dict, List, Any, Union, Optional
from pathlib import Path
import json
import joblib
import pandas as pd
import nltk
from transformers import AutoModel, AutoTokenizer
import torch
import numpy as np
from sklearn.base import TransformerMixin
class SimcseGenerator(TransformerMixin):
def __init__(
self, device: str ='cpu', batch_size: int =16, model_name: str = "princeton-nlp/unsup-simcse-bert-base-uncased"
) -> None:
self.model_name = model_name
self.device = torch.device(device)
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name).to(self.device)
self.tokenizer = tokenizer
self.model = model
self.batch_size = batch_size
def transform(self, X: np.ndarray) -> np.ndarray:
batch_size = (
16 # any larger, and we risk running out of memory on EC2 dev instances
)
embeddings = []
for start in range(0, len(X), batch_size):
end = min(len(X), start + batch_size)
inputs = self.tokenizer(
X[start:end],
padding=True,
truncation=True,
return_tensors="pt",
)
with torch.no_grad():
inputs = inputs.to(self.device)
batch_embeddings = self.model(
**inputs, output_hidden_states=True, return_dict=True
).pooler_output
embeddings.append(batch_embeddings.cpu().detach().numpy())
embeddings = np.concatenate(embeddings)
embeddings /= np.sqrt(np.square(embeddings).sum(axis=1))[:,np.newaxis]
return embeddings
class EndpointHandler():
def __init__(self, device: str = "cpu"):
# Preload all the elements you are going to need at inference.
# pseudo:
# self.model= load_model(path)
local_path = Path(__file__).parent
self.device = device
with open(local_path/'stop_words.json','r') as fp:
self.stop_words = set(json.load(fp))
with open(local_path/'instruction_label_map.json','r') as fp:
self.instruction_label_map = json.load(fp)
self.instruction_label_map = {int(k):v for k,v in self.instruction_label_map.items()}
self.instruction_pipeline = joblib.load(local_path/'instruction_classification_pipeline.joblib')
self.response_pipeline = joblib.load(local_path/'response_quality_pipeline.joblib')
self.simcse_generator = SimcseGenerator(device=self.device)
def _get_stop_word_proportion(self, s):
s = s.lower()
try:
words = nltk.tokenize.word_tokenize(s)
except:
words = nltk.tokenize.word_tokenize(s[1:])
if len(words)==0:
return 0
else:
return sum(x in self.stop_words for x in words) / len(words)
def predict_instruction_classes(self, df: pd.DataFrame) -> np.ndarray:
instruction_classes = self.instruction_pipeline.predict(df)
instruction_class_confidence = self.instruction_pipeline.predict_proba(df).max(axis=1)
return np.array(list(map(lambda x: self.instruction_label_map[x], instruction_classes))), instruction_class_confidence
def compute_response_quality_feature_space(self, df: pd.DataFrame, instruction_classes: Optional[np.ndarray] = None):
if instruction_classes is None:
instruction_classes, _ = self.predict_instruction_classes(df)
instruction_class_set = [self.instruction_label_map[i] for i in range(len(self.instruction_label_map))]
instruction_classes_onehot = pd.DataFrame(instruction_classes[:,np.newaxis]==np.array(instruction_class_set)[np.newaxis,:], columns=instruction_class_set).astype(float)
df1 = pd.concat([df,instruction_classes_onehot], axis=1)
df1['instruction_response_similarity'] = (self.simcse_generator.transform(df['instruction'].tolist()) * self.simcse_generator.transform(df['response'].tolist())).sum(axis=1)
df1['token_number'] = df1['response'].str.split().apply(len)
df1['stop_word_proportion'] = df1['response'].apply(self._get_stop_word_proportion)
return df1
def predict_response_quality(self, df, instruction_classes):
df1 = self.compute_response_quality_feature_space(df, instruction_classes)
return self.response_pipeline.predict_proba(df1)[:,1]
def __call__(self, df: Union[pd.DataFrame, Dict]):
is_dict = isinstance(df, dict)
if is_dict:
df = pd.DataFrame([df])
if 'dataset' not in df.columns:
df['dataset'] = ''
instruction_classes, instruction_class_confidences = self.predict_instruction_classes(df)
predictions = [{'instruction class': instruction_class, 'instruction class confidence': instruction_class_confidence} for instruction_class, instruction_class_confidence in zip(instruction_classes, instruction_class_confidences)]
if 'response' in df.columns:
response_qualities = self.predict_response_quality(df, instruction_classes)
for i,response_quality in enumerate(response_qualities):
predictions[i].update({'response quality': response_quality})
if is_dict:
return predictions[0]
else:
return pd.DataFrame(predictions, index=df.index) |