dream

1
import logging
2
import os
3
import time
4
from flask import Flask, request, jsonify
5
import sentry_sdk
6
from deeppavlov import build_model
7

8
logging.basicConfig(format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO)
9
logger = logging.getLogger(__name__)
10
sentry_sdk.init(os.getenv("SENTRY_DSN"))
11

12
app = Flask(__name__)
13

14
config_name = os.getenv("CONFIG")
15

16
try:
17
    el = build_model(config_name, download=True)
18
    logger.info("model loaded")
19
except Exception as e:
20
    sentry_sdk.capture_exception(e)
21
    logger.exception(e)
22
    raise e
23

24

25
@app.route("/model", methods=["POST"])
26
def respond():
27
    st_time = time.time()
28
    inp = request.json
29
    entity_substr_batch = inp.get("entity_substr", [[""]])
30
    entity_tags_batch = inp.get(
31
        "entity_tags", [["" for _ in entity_substr_list] for entity_substr_list in entity_substr_batch]
32
    )
33
    context_batch = inp.get("context", [[""]])
34
    opt_context_batch = []
35
    for entity_substr_list, hist_utt in zip(entity_substr_batch, context_batch):
36
        last_utt = hist_utt[-1]
37
        if last_utt[-1] not in {".", "!", "?"}:
38
            last_utt = f"{last_utt}."
39
        if len(hist_utt) > 1:
40
            prev_utt = hist_utt[-2]
41
            if prev_utt[-1] not in {".", "!", "?"}:
42
                prev_utt = f"{prev_utt}."
43
            opt_context_batch.append([prev_utt, last_utt])
44
        else:
45
            opt_context_batch.append([last_utt])
46

47
    entity_info_batch = [[{}] for _ in entity_substr_batch]
48
    try:
49
        entity_substr_batch, entity_ids_batch, conf_batch, entity_pages_batch = el(
50
            entity_substr_batch, entity_tags_batch, opt_context_batch
51
        )
52
        entity_info_batch = []
53
        for entity_substr_list, entity_ids_list, entity_tags_list, conf_list, entity_pages_list in zip(
54
            entity_substr_batch,
55
            entity_ids_batch,
56
            entity_tags_batch,
57
            conf_batch,
58
            entity_pages_batch,
59
        ):
60
            entity_info_list = []
61
            for entity_substr, entity_ids, entity_tags, confs, entity_pages in zip(
62
                entity_substr_list, entity_ids_list, entity_tags_list, conf_list, entity_pages_list
63
            ):
64
                entity_info = {}
65
                entity_info["entity_substr"] = entity_substr
66
                entity_info["entity_ids"] = entity_ids
67
                entity_info["entity_tags"] = entity_tags
68
                entity_info["confidences"] = [float(elem[2]) for elem in confs]
69
                entity_info["tokens_match_conf"] = [float(elem[0]) for elem in confs]
70
                entity_info["entity_pages"] = entity_pages
71
                entity_info_list.append(entity_info)
72
            entity_info_batch.append(entity_info_list)
73
    except Exception as e:
74
        sentry_sdk.capture_exception(e)
75
        logger.exception(e)
76
    total_time = time.time() - st_time
77
    logger.info(f"entity linking exec time = {total_time:.3f}s")
78
    return jsonify(entity_info_batch)
79

80

81
if __name__ == "__main__":
82
    app.run(debug=False, host="0.0.0.0", port=3000)
83