-
Notifications
You must be signed in to change notification settings - Fork 0
/
Game.py
78 lines (66 loc) · 2.34 KB
/
Game.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
import numpy as np
import torch
import torch.nn.functional as f
from transformers import BertModel, BertTokenizer
from time import sleep
from abc import ABC
class Analyzer(ABC):
def __init__(
self,
similarity_func,
bert_version : str,
available_words,
target
):
self.similarity_func = similarity_func
self.bert_version = bert_version
self.tokenizer = BertTokenizer.from_pretrained(bert_version)
self.model = BertModel.from_pretrained(bert_version)
self.model = self.model.eval()
self.observations= dict()
self.reward=0
self.done= False
self.action= " "
self.action_space= available_words
self.target = target
self.new_target= None
self.target_emb=self.get_word_emb(token=target) #np.random.choice(self.action_space)
self.attempts= 0
def get_word_emb(
self,
token : str
):
encoding = self.tokenizer(
token,
padding=True,
return_tensors='pt'
)
for tokens in encoding['input_ids']:
self.tokenizer.convert_ids_to_tokens(tokens)
with torch.no_grad():
embed = self.model(**encoding)[0]
avg_embed = embed.mean(dim=1)
return avg_embed
def act (self, word):
if not self.done:
if word not in self.action_space:
print("word not in avilable words")
self.reward= -1000
self.observations[word]= self.reward
self.attempts+= 1
return self.observations, self.reward, self.done, self.new_target
else:
self.reward= self.get_similarity(self.get_word_emb(token= word), self.target_emb)
self.observations[word]= self.reward
self.attempts+= 1
x= self.reward.item()
if (x >= 0.99999):
self.done= True
self.new_target= np.random.choice(self.action_space)
return self.observations, self.reward, self.done, self.new_target
elif self.done:
return self.observations, self.reward, self.done ,self.new_target
def get_similarity(self, embed1, embed2):
return self.similarity_func(embed1, embed2)[0]
def giveup(self):
return self.target