update handler
Browse files- handler.py +24 -116
handler.py
CHANGED
@@ -1,54 +1,15 @@
|
|
1 |
import os
|
|
|
2 |
import json
|
3 |
import torch
|
4 |
-
import numpy as np
|
5 |
-
from transformers import BertTokenizer
|
6 |
from ts.torch_handler.base_handler import BaseHandler
|
7 |
-
from
|
8 |
-
|
9 |
-
import transformers
|
10 |
-
import torch
|
11 |
-
import torch.nn as nn
|
12 |
-
import torch.nn.functional as F
|
13 |
-
|
14 |
-
class AttentionPool(nn.Module):
|
15 |
-
def __init__(self, hidden_size):
|
16 |
-
super().__init__()
|
17 |
-
self.attention = nn.Linear(hidden_size, 1)
|
18 |
-
|
19 |
-
def forward(self, last_hidden_state):
|
20 |
-
attention_scores = self.attention(last_hidden_state).squeeze(-1)
|
21 |
-
attention_weights = F.softmax(attention_scores, dim=1)
|
22 |
-
pooled_output = torch.bmm(attention_weights.unsqueeze(1), last_hidden_state).squeeze(1)
|
23 |
-
return pooled_output
|
24 |
-
|
25 |
-
class MultiSampleDropout(nn.Module):
|
26 |
-
def __init__(self, dropout=0.5, num_samples=5):
|
27 |
-
super().__init__()
|
28 |
-
self.dropout = nn.Dropout(dropout)
|
29 |
-
self.num_samples = num_samples
|
30 |
-
|
31 |
-
def forward(self, x):
|
32 |
-
return torch.mean(torch.stack([self.dropout(x) for _ in range(self.num_samples)]), dim=0)
|
33 |
-
|
34 |
|
35 |
-
|
36 |
-
|
37 |
-
|
38 |
-
self.bert = transformers.BertModel.from_pretrained('bert-base-uncased')
|
39 |
-
self.attention_pool = AttentionPool(768)
|
40 |
-
self.dropout = MultiSampleDropout()
|
41 |
-
self.norm = nn.LayerNorm(768)
|
42 |
-
self.classifier = nn.Linear(768, num_classes)
|
43 |
-
|
44 |
-
def forward(self, input_ids, attention_mask, token_type_ids):
|
45 |
-
bert_output = self.bert(input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
|
46 |
-
pooled_output = self.attention_pool(bert_output.last_hidden_state)
|
47 |
-
pooled_output = self.dropout(pooled_output)
|
48 |
-
pooled_output = self.norm(pooled_output)
|
49 |
-
logits = self.classifier(pooled_output)
|
50 |
-
return logits
|
51 |
|
|
|
52 |
|
53 |
class UICardMappingHandler(BaseHandler):
|
54 |
def __init__(self):
|
@@ -59,96 +20,43 @@ class UICardMappingHandler(BaseHandler):
|
|
59 |
self.manifest = context.manifest
|
60 |
properties = context.system_properties
|
61 |
model_dir = properties.get("model_dir")
|
62 |
-
|
63 |
-
|
64 |
-
|
65 |
-
# Load config
|
66 |
-
with open(os.path.join(model_dir, 'config.json'), 'r') as f:
|
67 |
-
self.config = json.load(f)
|
68 |
-
|
69 |
-
# Initialize encoder and labels
|
70 |
-
self.labels = ['Videos', 'Unit Conversion', 'Translation', 'Shopping Product Comparison', 'Restaurants', 'Product', 'Information', 'Images', 'Gift', 'General Comparison', 'Flights', 'Answer', 'Aircraft Seat Map']
|
71 |
-
labels_np = np.array(self.labels).reshape(-1, 1)
|
72 |
-
self.encoder = OneHotEncoder(sparse_output=False)
|
73 |
-
self.encoder.fit(labels_np)
|
74 |
-
|
75 |
-
# Load model
|
76 |
self.model = ImprovedBERTClass()
|
77 |
self.model.load_state_dict(torch.load(os.path.join(model_dir, 'model.pth'), map_location=self.device))
|
78 |
self.model.to(self.device)
|
79 |
self.model.eval()
|
80 |
-
|
81 |
-
# Load tokenizer
|
82 |
-
self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
|
83 |
-
|
84 |
self.initialized = True
|
85 |
|
86 |
def preprocess(self, data):
|
87 |
-
text = data[0].get("
|
88 |
-
|
89 |
-
|
90 |
-
inputs = self.tokenizer
|
91 |
-
|
92 |
-
add_special_tokens=True,
|
93 |
-
max_length=64,
|
94 |
-
padding='max_length',
|
95 |
-
return_tensors='pt',
|
96 |
-
truncation=True
|
97 |
-
)
|
98 |
-
|
99 |
-
return {
|
100 |
-
"ids": inputs['input_ids'].to(self.device, dtype=torch.long),
|
101 |
-
"mask": inputs['attention_mask'].to(self.device, dtype=torch.long),
|
102 |
-
"token_type_ids": inputs['token_type_ids'].to(self.device, dtype=torch.long),
|
103 |
-
"k": k
|
104 |
-
}
|
105 |
|
106 |
-
def inference(self,
|
107 |
with torch.no_grad():
|
108 |
-
outputs = self.model(
|
109 |
-
|
110 |
-
return probabilities.cpu().detach().numpy().flatten(), data["k"]
|
111 |
|
112 |
def postprocess(self, inference_output):
|
113 |
-
probabilities
|
|
|
114 |
|
115 |
-
#
|
116 |
-
top_k_indices =
|
117 |
top_k_probs = probabilities[top_k_indices]
|
118 |
|
119 |
-
|
120 |
-
top_k_one_hot = np.zeros((k, len(probabilities)))
|
121 |
-
for i, idx in enumerate(top_k_indices):
|
122 |
-
top_k_one_hot[i, idx] = 1
|
123 |
-
|
124 |
-
# Decode the top k predictions
|
125 |
-
top_k_cards = [self.decode_vector(one_hot.reshape(1, -1)) for one_hot in top_k_one_hot]
|
126 |
|
127 |
-
|
128 |
-
top_k_predictions = list(zip(top_k_cards, top_k_probs.tolist()))
|
129 |
|
130 |
-
# Determine the most likely card
|
131 |
-
predicted_labels = (probabilities > 0.5).astype(int)
|
132 |
-
if sum(predicted_labels) == 0:
|
133 |
-
most_likely_card = "Answer"
|
134 |
-
else:
|
135 |
-
most_likely_card = self.decode_vector(predicted_labels.reshape(1, -1))
|
136 |
-
|
137 |
-
# Prepare the response
|
138 |
result = {
|
139 |
"most_likely_card": most_likely_card,
|
140 |
"top_k_predictions": top_k_predictions
|
141 |
}
|
142 |
|
143 |
return [result]
|
144 |
-
|
145 |
-
def decode_vector(self, vector):
|
146 |
-
original_label = self.encoder.inverse_transform(vector)
|
147 |
-
return original_label[0][0] # Returns the label as a string
|
148 |
-
|
149 |
-
def handle(self, data, context):
|
150 |
-
self.context = context
|
151 |
-
data = self.preprocess(data)
|
152 |
-
data = self.inference(data)
|
153 |
-
data = self.postprocess(data)
|
154 |
-
return data
|
|
|
1 |
import os
|
2 |
+
import sys
|
3 |
import json
|
4 |
import torch
|
|
|
|
|
5 |
from ts.torch_handler.base_handler import BaseHandler
|
6 |
+
from transformers import BertTokenizer
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
7 |
|
8 |
+
# Add the model directory to the Python path
|
9 |
+
model_dir = os.path.dirname(os.path.abspath(__file__))
|
10 |
+
sys.path.append(model_dir)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
11 |
|
12 |
+
from model import ImprovedBERTClass # Ensure this import matches your model file name
|
13 |
|
14 |
class UICardMappingHandler(BaseHandler):
|
15 |
def __init__(self):
|
|
|
20 |
self.manifest = context.manifest
|
21 |
properties = context.system_properties
|
22 |
model_dir = properties.get("model_dir")
|
23 |
+
self.device = torch.device("cuda:" + str(properties.get("gpu_id")) if torch.cuda.is_available() else "cpu")
|
24 |
+
|
25 |
+
self.tokenizer = BertTokenizer.from_pretrained(model_dir)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
26 |
self.model = ImprovedBERTClass()
|
27 |
self.model.load_state_dict(torch.load(os.path.join(model_dir, 'model.pth'), map_location=self.device))
|
28 |
self.model.to(self.device)
|
29 |
self.model.eval()
|
30 |
+
|
|
|
|
|
|
|
31 |
self.initialized = True
|
32 |
|
33 |
def preprocess(self, data):
|
34 |
+
text = data[0].get("data")
|
35 |
+
if text is None:
|
36 |
+
text = data[0].get("body")
|
37 |
+
inputs = self.tokenizer(text, return_tensors="pt", max_length=64, padding='max_length', truncation=True)
|
38 |
+
return inputs.to(self.device)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
39 |
|
40 |
+
def inference(self, inputs):
|
41 |
with torch.no_grad():
|
42 |
+
outputs = self.model(**inputs)
|
43 |
+
return torch.sigmoid(outputs.logits)
|
|
|
44 |
|
45 |
def postprocess(self, inference_output):
|
46 |
+
probabilities = inference_output.cpu().numpy().flatten()
|
47 |
+
labels = ['Videos', 'Unit Conversion', 'Translation', 'Shopping Product Comparison', 'Restaurants', 'Product', 'Information', 'Images', 'Gift', 'General Comparison', 'Flights', 'Answer', 'Aircraft Seat Map']
|
48 |
|
49 |
+
top_k = 3 # You can adjust this value
|
50 |
+
top_k_indices = probabilities.argsort()[-top_k:][::-1]
|
51 |
top_k_probs = probabilities[top_k_indices]
|
52 |
|
53 |
+
top_k_predictions = [{"card": labels[i], "probability": float(p)} for i, p in zip(top_k_indices, top_k_probs)]
|
|
|
|
|
|
|
|
|
|
|
|
|
54 |
|
55 |
+
most_likely_card = "Answer" if sum(probabilities > 0.5) == 0 else labels[probabilities.argmax()]
|
|
|
56 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
57 |
result = {
|
58 |
"most_likely_card": most_likely_card,
|
59 |
"top_k_predictions": top_k_predictions
|
60 |
}
|
61 |
|
62 |
return [result]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|