forked from JinkyungJo/KoBART_weather
-
Notifications
You must be signed in to change notification settings - Fork 0
/
infer_dev.py
196 lines (177 loc) · 6.44 KB
/
infer_dev.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
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
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
import torch
from kobart import get_kobart_tokenizer
from transformers.models.bart import BartForConditionalGeneration
from sentence_transformers import SentenceTransformer, util
from tokenizers import Tokenizer
from grammar_regex import is_correct_grammar
import pandas as pd
from random import randrange
from pprint import pprint
import numpy as np
def load_model():
model = BartForConditionalGeneration.from_pretrained('./kobart_weather_v2')
model2 = SentenceTransformer('sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2')
return model, model2
def get_tokenizer():
return get_kobart_tokenizer()
def get_sql(input, templates):
global model2
template_embeds = templates[0]
index_to_input = templates[1]
template_dict = templates[2]
embeds = model2.encode(input)
#Compute cosine-similarities for input and input templates for matching
cosine_scores = util.pytorch_cos_sim(embeds, template_embeds)
indx = np.argmax((cosine_scores.numpy())[0])
input_template = index_to_input[indx]
output_template = template_dict[input_template] #matched sql template
checkpoints = []
j=0
# Getting checkpoints between matched template & input
for i in range(len(input_template)):
if input_template[i] in input[j:]:
j_ = input[j:].index(input_template[i])
j = j + j_
checkpoints.append([i,j])
ot = output_template
#Iterate through the checkpoints
for i in range(len(checkpoints)-1):
t1 = checkpoints[i][0]
o1 = checkpoints[i][1]
t2 = checkpoints[i+1][0]
o2 = checkpoints[i+1][1]
if (t1+1)!=t2:
template_var = input_template[t1+1:t2]
input_var = input[o1+1:o2]
print(input_var)
if template_var in ot:
if 'date' in template_var: #handling date
year_ = input_var.split('년')
year = year_[0]
month_ = year_[1].split('월')
month = month_[0][1:]
if len(month)==1:
month = f'0{month}'
day_ = month_[1].split('일')
day = day_[0][1:]
if len(day)==1:
day = f'0{day}'
input_var = f"'{year+month+day}'"
ot = ot.replace(template_var, input_var)
elif 'month' in template_var:
month= (input_var.split('월'))[0]
input_var = f"'{month}'"
ot = ot.replace(template_var, input_var)
elif 'number' in template_var:
ot = ot.replace(template_var, input_var)
else:
input_var = f"'{input_var}'"
ot = ot.replace(template_var, input_var)
if ot == output_template:
return []
else:
return ot
def get_output(input, templates):
text = input['source']
date_s = input['date'].split(" ")
ymd = date_s[0].split('-')
hms = date_s[1].split(':')
date = [ymd[0],ymd[1],ymd[2],hms[0],hms[1]]
# Get rid of date information input
original_text = text
if '-' in text and ':' in text:
text = text[17:]
elif '-' in text:
text = text[11:]
elif ':' in text:
text = text[6:]
input_ids = tokenizer.encode(text)
input_ids = torch.tensor(input_ids).to('cuda')
input_ids = input_ids.unsqueeze(0)
outputs = model.generate(input_ids, eos_token_id=1, max_length=512, num_beams=5, num_return_sequences=5)
res = []
for output in outputs:
res.append(tokenizer.decode(output, skip_special_tokens=True))
out = None
for x in res:
if is_correct_grammar(x):
criteria_met = True
out = x
break
if out==None:
out = res[0]
sql = get_sql(text, templates)
return [original_text, out, date, sql]
def response_template(res):
input = res[0]
output = res[1]
date = res[2]
sql = res[3]
if date!=[]:
year = date[0]
month = date[1]
day = date[2]
hour = date[3]
minute = date[4]
if '내일' in input:
day = str(int(day) + 1)
if '어제' in input:
day = str(int(day) - 1)
#handling custom year-month-day time"
# There is date information included in input
if '-' in input:
indx = input.find('-')
year = input[indx-4:indx]
month = input[indx+1:indx+3]
day = input[indx+4:indx+6]
# There is time information included in input
if ':' in input:
indx = input.find(':')
hour = input[indx-2:indx]
minute = input[indx+1:indx+3]
output = output.replace('YYYYMMDDHHMI', year+month+day+hour+minute)
else:
output = output.replace("입력='YYYYMMDDHHMI'", '')
response = {
"pseudoList":[{
"site":"COMIS",
"pseudo":output,
}, {}],
"extremeValue":[sql]
}
return response
def get_template_embeddings(model, data_dir):
try:
sql_template = pd.read_csv("data/template.csv")
except:
sql_template = pd.read_csv('home/KoBART-summarization/template.csv')
template_dict = {}
index_to_input = {}
template_embeds = []
# Getting templates
for index,row in sql_template.iterrows():
input = row['input']
output = row['output']
index_to_input[index] = input
template_dict[input] = output
template_embeds.append(model.encode(input))
return (template_embeds, index_to_input, template_dict)
tokenizer = get_tokenizer()
if __name__ == '__main__':
#example = "전기간 전지점 일단위 최고온도 3개"
example = "당일(2021년 1월 2일) 전지점 일단위 최고온도 30개"
#example = "당일(2021년 1월 1일) 전지점 일단위 최저온도 3개"
input = {
"source" : example,
"date" : "2021-11-16 00:00:00",
"sourceType" : "text",
"responseChannel": "aiw-response"
}
model, model2 = load_model()
model = model.to('cuda')
model2 = model2.to('cuda')
tokenizer = get_tokenizer()
templates = get_template_embeddings(model2, 'template.csv')
output = response_template(get_output(input, templates))
print('input:', input)
print('output:', output)