-
Notifications
You must be signed in to change notification settings - Fork 103
/
app.py
73 lines (57 loc) · 1.98 KB
/
app.py
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
import config
import torch
import flask
import time
from flask import Flask
from flask import request
from model import BERTBaseUncased
import functools
import torch.nn as nn
app = Flask(__name__)
MODEL = None
DEVICE = config.DEVICE
PREDICTION_DICT = dict()
def sentence_prediction(sentence):
tokenizer = config.TOKENIZER
max_len = config.MAX_LEN
review = str(sentence)
review = " ".join(review.split())
inputs = tokenizer.encode_plus(
review, None, add_special_tokens=True, max_length=max_len
)
ids = inputs["input_ids"]
mask = inputs["attention_mask"]
token_type_ids = inputs["token_type_ids"]
padding_length = max_len - len(ids)
ids = ids + ([0] * padding_length)
mask = mask + ([0] * padding_length)
token_type_ids = token_type_ids + ([0] * padding_length)
ids = torch.tensor(ids, dtype=torch.long).unsqueeze(0)
mask = torch.tensor(mask, dtype=torch.long).unsqueeze(0)
token_type_ids = torch.tensor(token_type_ids, dtype=torch.long).unsqueeze(0)
ids = ids.to(DEVICE, dtype=torch.long)
token_type_ids = token_type_ids.to(DEVICE, dtype=torch.long)
mask = mask.to(DEVICE, dtype=torch.long)
outputs = MODEL(ids=ids, mask=mask, token_type_ids=token_type_ids)
outputs = torch.sigmoid(outputs).cpu().detach().numpy()
return outputs[0][0]
@app.route("/predict")
def predict():
sentence = request.args.get("sentence")
start_time = time.time()
positive_prediction = sentence_prediction(sentence)
negative_prediction = 1 - positive_prediction
response = {}
response["response"] = {
"positive": str(positive_prediction),
"negative": str(negative_prediction),
"sentence": str(sentence),
"time_taken": str(time.time() - start_time),
}
return flask.jsonify(response)
if __name__ == "__main__":
MODEL = BERTBaseUncased()
MODEL.load_state_dict(torch.load(config.MODEL_PATH))
MODEL.to(DEVICE)
MODEL.eval()
app.run(host="0.0.0.0", port="9999")