Spaces:
Running
on
CPU Upgrade
Running
on
CPU Upgrade
File size: 4,601 Bytes
63858e7 |
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 175 176 177 178 |
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 data_processing import from_model
from transformer_details import from_pretrained
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=5050, 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):
model = request['model']
deets = from_pretrained(model)
info = deets.model.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_attention_and_meta(**request):
model = request["model"]
details = from_pretrained(model)
sentence = request["sentence"]
layer = int(request["layer"])
deets = details.att_from_sentence(sentence)
payload_out = deets.to_json(layer)
return {
"status": 200,
"payload": payload_out
}
def update_masked_attention(**request):
"""
Return attention information from tokens and mask indices.
Object: {"a" : {"sentence":__, "mask_inds"}, "b" : {...}}
"""
payload = request["payload"]
model = payload['model']
details = from_pretrained(model)
tokens = payload["tokens"]
sentence = payload["sentence"]
mask = payload["mask"]
layer = int(payload["layer"])
MASK = details.aligner.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.att_from_tokens(token_inputs, sentence)
payload_out = deets.to_json(layer)
return {
"status": 200,
"payload": payload_out,
}
def nearest_embedding_search(**request):
"""Return the token text and the metadata in JSON"""
model = request["model"]
corpus = request["corpus"]
try:
details = from_pretrained(model)
except KeyError as e:
return {'status': 405, "payload": None}
try:
cc = from_model(model, corpus)
except FileNotFoundError as e:
return {
"status": 406,
"payload": None
}
q = np.array(request["embedding"]).reshape((1, -1)).astype(np.float32)
layer = int(request["layer"])
heads = list(map(int, list(set(request["heads"]))))
k = int(request["k"])
out = cc.search_embeddings(layer, q, k)
payload_out = [o.to_json(layer, heads) for o in out]
return {
"status": 200,
"payload": payload_out
}
def nearest_context_search(**request):
"""Return the token text and the metadata in JSON"""
model = request["model"]
corpus = request["corpus"]
print("CORPUS: ", corpus)
try:
details = from_pretrained(model)
except KeyError as e:
return {'status': 405, "payload": None}
try:
cc = from_model(model, corpus)
except FileNotFoundError as e:
return {'status': 406, "payload": None}
q = np.array(request["context"]).reshape((1, -1)).astype(np.float32)
layer = int(request["layer"])
heads = list(map(int, list(set(request["heads"]))))
k = int(request["k"])
out = cc.search_contexts(layer, heads, q, k)
payload_out = [o.to_json(layer, heads) for o in out]
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)
|