File size: 5,754 Bytes
a94233d e732be5 a94233d e732be5 5aa95d8 e732be5 a94233d e732be5 a94233d e732be5 a94233d e732be5 a94233d e732be5 a94233d e732be5 a94233d e732be5 a94233d e732be5 a94233d e732be5 a94233d e732be5 a94233d e732be5 a94233d e732be5 |
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 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 |
import os
import json
import torch
import numpy as np
from transformers import BertTokenizer
from ts.torch_handler.base_handler import BaseHandler
from sklearn.preprocessing import OneHotEncoder
import transformers
import torch
import torch.nn as nn
import torch.nn.functional as F
class AttentionPool(nn.Module):
def __init__(self, hidden_size):
super().__init__()
self.attention = nn.Linear(hidden_size, 1)
def forward(self, last_hidden_state):
attention_scores = self.attention(last_hidden_state).squeeze(-1)
attention_weights = F.softmax(attention_scores, dim=1)
pooled_output = torch.bmm(attention_weights.unsqueeze(1), last_hidden_state).squeeze(1)
return pooled_output
class MultiSampleDropout(nn.Module):
def __init__(self, dropout=0.5, num_samples=5):
super().__init__()
self.dropout = nn.Dropout(dropout)
self.num_samples = num_samples
def forward(self, x):
return torch.mean(torch.stack([self.dropout(x) for _ in range(self.num_samples)]), dim=0)
class ImprovedBERTClass(nn.Module):
def __init__(self, num_classes=13):
super().__init__()
self.bert = transformers.BertModel.from_pretrained('bert-base-uncased')
self.attention_pool = AttentionPool(768)
self.dropout = MultiSampleDropout()
self.norm = nn.LayerNorm(768)
self.classifier = nn.Linear(768, num_classes)
def forward(self, input_ids, attention_mask, token_type_ids):
bert_output = self.bert(input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
pooled_output = self.attention_pool(bert_output.last_hidden_state)
pooled_output = self.dropout(pooled_output)
pooled_output = self.norm(pooled_output)
logits = self.classifier(pooled_output)
return logits
class UICardMappingHandler(BaseHandler):
def __init__(self):
super().__init__()
self.initialized = False
def initialize(self, context):
self.manifest = context.manifest
properties = context.system_properties
model_dir = properties.get("model_dir")
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Load config
with open(os.path.join(model_dir, 'config.json'), 'r') as f:
self.config = json.load(f)
# Initialize encoder and labels
self.labels = ['Videos', 'Unit Conversion', 'Translation', 'Shopping Product Comparison', 'Restaurants', 'Product', 'Information', 'Images', 'Gift', 'General Comparison', 'Flights', 'Answer', 'Aircraft Seat Map']
labels_np = np.array(self.labels).reshape(-1, 1)
self.encoder = OneHotEncoder(sparse_output=False)
self.encoder.fit(labels_np)
# Load model
self.model = ImprovedBERTClass()
self.model.load_state_dict(torch.load(os.path.join(model_dir, 'model.pth'), map_location=self.device))
self.model.to(self.device)
self.model.eval()
# Load tokenizer
self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
self.initialized = True
def preprocess(self, data):
text = data[0].get("body").get("text", "")
k = data[0].get("body").get("k", 3)
inputs = self.tokenizer.encode_plus(
text,
add_special_tokens=True,
max_length=64,
padding='max_length',
return_tensors='pt',
truncation=True
)
return {
"ids": inputs['input_ids'].to(self.device, dtype=torch.long),
"mask": inputs['attention_mask'].to(self.device, dtype=torch.long),
"token_type_ids": inputs['token_type_ids'].to(self.device, dtype=torch.long),
"k": k
}
def inference(self, data):
with torch.no_grad():
outputs = self.model(data["ids"], data["mask"], data["token_type_ids"])
probabilities = torch.sigmoid(outputs)
return probabilities.cpu().detach().numpy().flatten(), data["k"]
def postprocess(self, inference_output):
probabilities, k = inference_output
# Get top k predictions
top_k_indices = np.argsort(probabilities)[-k:][::-1]
top_k_probs = probabilities[top_k_indices]
# Create one-hot encodings for top k indices
top_k_one_hot = np.zeros((k, len(probabilities)))
for i, idx in enumerate(top_k_indices):
top_k_one_hot[i, idx] = 1
# Decode the top k predictions
top_k_cards = [self.decode_vector(one_hot.reshape(1, -1)) for one_hot in top_k_one_hot]
# Create a list of tuples (card, probability) for top k predictions
top_k_predictions = list(zip(top_k_cards, top_k_probs.tolist()))
# Determine the most likely card
predicted_labels = (probabilities > 0.5).astype(int)
if sum(predicted_labels) == 0:
most_likely_card = "Answer"
else:
most_likely_card = self.decode_vector(predicted_labels.reshape(1, -1))
# Prepare the response
result = {
"most_likely_card": most_likely_card,
"top_k_predictions": top_k_predictions
}
return [result]
def decode_vector(self, vector):
original_label = self.encoder.inverse_transform(vector)
return original_label[0][0] # Returns the label as a string
def handle(self, data, context):
self.context = context
data = self.preprocess(data)
data = self.inference(data)
data = self.postprocess(data)
return data
|