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
top_n = int(os.getenv("TOP_N"))
16

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

25

26
@app.route("/model", methods=["POST"])
27
def respond():
28
    st_time = time.time()
29
    inp = request.json
30
    dialog_history_batch = inp.get("dialog_history", [])
31
    entity_substr_batch = inp.get("entity_substr", [[] for _ in dialog_history_batch])
32
    entity_tags_batch = inp.get("entity_tags", [[] for _ in dialog_history_batch])
33
    entity_pages_batch = inp.get("entity_pages", [[] for _ in dialog_history_batch])
34
    sentences_batch = []
35
    for dialog_history in dialog_history_batch:
36
        if (len(dialog_history[-1].split()) > 2 and "?" in dialog_history[-1]) or len(dialog_history) == 1:
37
            sentence = dialog_history[-1]
38
        else:
39
            sentence = " ".join(dialog_history)
40
        sentences_batch.append(sentence)
41

42
    contexts_with_scores_batch = [[] for _ in sentences_batch]
43
    try:
44
        contexts_with_scores_batch = []
45
        contexts_batch, scores_batch, from_linked_page_batch, numbers_batch = fact_retrieval(
46
            sentences_batch, entity_substr_batch, entity_tags_batch, entity_pages_batch
47
        )
48
        for contexts, scores, from_linked_page_list, numbers in zip(
49
            contexts_batch, scores_batch, from_linked_page_batch, numbers_batch
50
        ):
51
            contexts_with_scores_linked, contexts_with_scores_not_linked, contexts_with_scores_first = [], [], []
52
            for context, score, from_linked_page, number in zip(contexts, scores, from_linked_page_list, numbers):
53
                if from_linked_page and number > 0:
54
                    contexts_with_scores_linked.append((context, score, number))
55
                elif from_linked_page and number == 0:
56
                    contexts_with_scores_first.append((context, score, number))
57
                else:
58
                    contexts_with_scores_not_linked.append((context, score, number))
59
            contexts_with_scores_linked = sorted(contexts_with_scores_linked, key=lambda x: (x[1], x[2]), reverse=True)
60
            contexts_with_scores_not_linked = sorted(
61
                contexts_with_scores_not_linked, key=lambda x: (x[1], x[2]), reverse=True
62
            )
63
            contexts_with_scores = []
64
            contexts_with_scores += [(context, score, True) for context, score, _ in contexts_with_scores_first]
65
            contexts_with_scores += [
66
                (context, score, True) for context, score, _ in contexts_with_scores_linked[: top_n // 2]
67
            ]
68
            contexts_with_scores += [
69
                (context, score, False) for context, score, _ in contexts_with_scores_not_linked[: top_n // 2]
70
            ]
71
            contexts_with_scores_batch.append(contexts_with_scores)
72
    except Exception as e:
73
        sentry_sdk.capture_exception(e)
74
        logger.exception(e)
75
    total_time = time.time() - st_time
76
    logger.info(f"fact retrieval exec time = {total_time:.3f}s")
77
    return jsonify(contexts_with_scores_batch)
78

79

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