-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathmodels.py
107 lines (77 loc) · 3.09 KB
/
models.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
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
from typing import List
class Entity:
ID = None # type: str
full_title = None # type: str
title = None # type: str
sub_title = None # type: str
language = None # type: str
source = None # type: str
embed = None # type: List[float]
def __init__(self, entity_id, title, sub_title, source, language, embed=None):
self.ID = entity_id
self.full_title = title + sub_title
self.title = title
self.sub_title = sub_title
self.source = source
self.language = language
self.embed = embed
def set_embed(self, embed: List[float]):
self.embed = embed
def get_full_title(self):
return self.full_title
class Candidate:
entity_id = None # type: str
entity = None # type: Entity
context_words_sim = None # type: float
context_entities_sim = None # type: float
e_given_m = None # type: float
believe_score = None # type: float
def __init__(self, entity_id, entity_title=""):
self.entity_id = entity_id
self.entity_title = entity_title
def set_context_words_sim(self, similarity: float):
self.context_words_sim = similarity
def set_context_entities_sim(self, similarity: float):
self.context_entities_sim = similarity
def set_e_given_m(self, prob: float):
self.e_given_m = prob
def set_entity(self, entity: Entity):
self.entity = entity
def set_believe_score(self, score: float):
self.believe_score = score
class Mention:
label = None # type: str
candidates = [] # type: List[Candidate]
start = None # type: int
end = None # type: int
prev_context = None # type: List[str]
after_context = None # type: List[str]
context_entities = None # type: List[Entity]
result_cand = None # type: Candidate
gold_entity = None # type: Entity
believe_score = None # type: float
parse_from = None # type: str
def __init__(self, start, end, mention_str, candidates=None):
self.start = start
self.end = end
self.label = mention_str
if candidates is None:
self.candidates = []
else:
self.candidates = Candidate
def add_candidate(self, candidate: Candidate):
self.candidates.append(candidate)
def set_prev_context(self, context_words):
self.prev_context = context_words
def set_after_context(self, context_words):
self.after_context = context_words
def set_context_entities(self, context_entities: List[Entity]):
self.context_entities = context_entities
def set_result_cand(self, candidate):
self.result_cand = candidate
def set_gold_entity(self, entity: Entity):
self.gold_entity = entity
def set_believe_score(self, score: float):
self.believe_score = score
def __str__(self):
return "{}, {}, {}, {}".format(self.start, self.end, self.label, "::=".join([cand.entity_id for cand in self.candidates]))