Spaces:
Sleeping
Sleeping
File size: 4,900 Bytes
df92393 |
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 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 |
import argparse
import numpy as np
import connexion
from flask_cors import CORS
from flask import render_template, redirect, send_from_directory
import utils.path_fixes as pf
from utils.f import ifnone
from model_api import get_details
app = connexion.FlaskApp(__name__, static_folder="client/dist", specification_dir=".")
flask_app = app.app
CORS(flask_app)
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument("--debug", action="store_true", help=" Debug mode")
parser.add_argument("--port", default=5051, help="Port to run the app. ")
# Flask main routes
@app.route("/")
def hello_world():
return redirect("client/exBERT.html")
# send everything from client as static content
@app.route("/client/<path:path>")
def send_static_client(path):
""" serves all files from ./client/ to ``/client/<path:path>``
:param path: path from api call
"""
return send_from_directory(str(pf.CLIENT_DIST), path)
# ======================================================================
## CONNEXION API ##
# ======================================================================
def get_model_details(**request):
"""Get important information about a model, like the number of layers and heads
Args:
request['model']: The model name
Returns:
{
status: 200,
payload: {
nlayers (int)
nheads (int)
}
}
"""
mname = request['model']
deets = get_details(mname)
info = deets.config
nlayers = info.num_hidden_layers
nheads = info.num_attention_heads
payload_out = {
"nlayers": nlayers,
"nheads": nheads,
}
return {
"status": 200,
"payload": payload_out,
}
def get_attentions_and_preds(**request):
"""For a sentence, at a layer, get the attentions and predictions
Args:
request['model']: Model name
request['sentence']: Sentence to get the attentions for
request['layer']: Which layer to extract from
Returns:
{
status: 200
payload: {
aa: {
att: Array((nheads, ntoks, ntoks))
left: [{
text (str),
topk_words (List[str]),
topk_probs (List[float])
}, ...]
right: [{
text (str),
topk_words (List[str]),
topk_probs (List[float])
}, ...]
}
}
}
"""
model = request["model"]
details = get_details(model)
sentence = request["sentence"]
layer = int(request["layer"])
deets = details.from_sentence(sentence)
payload_out = deets.to_json(layer)
return {
"status": 200,
"payload": payload_out
}
def update_masked_attention(**request):
"""From tokens and indices of what should be masked, get the attentions and predictions
payload = request['payload']
Args:
payload['model'] (str): Model name
payload['tokens'] (List[str]): Tokens to pass through the model
payload['sentence'] (str): Original sentence the tokens came from
payload['mask'] (List[int]): Which indices to mask
payload['layer'] (int): Which layer to extract information from
Returns:
{
status: 200
payload: {
aa: {
att: Array((nheads, ntoks, ntoks))
left: [{
text (str),
topk_words (List[str]),
topk_probs (List[float])
}, ...]
right: [{
text (str),
topk_words (List[str]),
topk_probs (List[float])
}, ...]
}
}
}
"""
payload = request["payload"]
model = payload['model']
details = get_details(model)
tokens = payload["tokens"]
sentence = payload["sentence"]
mask = payload["mask"]
layer = int(payload["layer"])
MASK = details.tok.mask_token
mask_tokens = lambda toks, maskinds: [
t if i not in maskinds else ifnone(MASK, t) for (i, t) in enumerate(toks)
]
token_inputs = mask_tokens(tokens, mask)
deets = details.from_tokens(token_inputs, sentence)
payload_out = deets.to_json(layer)
return {
"status": 200,
"payload": payload_out,
}
app.add_api("swagger.yaml")
# Setup code
if __name__ != "__main__":
print("SETTING UP ENDPOINTS")
# Then deploy app
else:
args, _ = parser.parse_known_args()
print("Initiating app")
app.run(port=args.port, use_reloader=False, debug=args.debug) |