-
Notifications
You must be signed in to change notification settings - Fork 0
/
rg_models.py
65 lines (49 loc) · 2.16 KB
/
rg_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
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM
import torch
device = "cuda" if torch.cuda.is_available() else "cpu"
class RGModel():
def __init__(self):
pass
def predict(self, history):
raise NotImplementedError()
class GelatoRGModel(RGModel):
def __init__(self):
super().__init__()
self.model_path = "./output/rg_model/checkpoint-best"
self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)
self.model = AutoModelForSeq2SeqLM.from_pretrained(self.model_path).to(device)
def history_to_string(self, history):
assert isinstance(history, list)
processed_history = " ".join(list(map(lambda x: x["speaker"] + ": " + x["utterance"], history)))
return processed_history
def predict(self, history):
prefix = "response generation"
context_text = self.history_to_string(history)
inputs = prefix + " : " + context_text
model_inputs = self.tokenizer([inputs], return_tensors="pt").to(device)
generated_ids = self.model.generate(**model_inputs, max_new_tokens=512)
output = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)[0]
return output
if __name__ == '__main__':
rg_model = GelatoRGModel()
history = [{
"speaker": "customer",
"utterance": "Hello! Could you please tell me if there is a vegan option available today?"
},
{
"speaker": "assistant",
"utterance": "Yes, we have several vegan options: Gianduja, Passion Fruit Sorbet, Dark Chocolate & Sea Salt, and Coconut, Raspberry Ripple."
},
{
"speaker": "customer",
"utterance": "I'd like a double scoop with Passion Fruit Sorbet and Dark Chocolate & Sea Salt, in a normal cone.",
"state": {
"flavours": [
"Passion Fruit Sorbet",
"Dark Chocolate & Sea Salt"
],
"size": "Double Scoop",
"container": "Normal Cone"
}
}]
print(rg_model.predict(history))