Spaces:
Runtime error
Runtime error
AmitGarage
commited on
Commit
•
d2ef806
1
Parent(s):
a25ba4b
Update scripts/torch_ner_model.py
Browse files- scripts/torch_ner_model.py +8 -42
scripts/torch_ner_model.py
CHANGED
@@ -21,6 +21,7 @@ def build_torch_ner_model(
|
|
21 |
hidden_width: int,
|
22 |
dropout: Optional[float] = None,
|
23 |
nO: Optional[int] = None,
|
|
|
24 |
) -> Model[List[Doc], List[Floats2d]]:
|
25 |
"""Build a tagger model, using a provided token-to-vector component. The tagger
|
26 |
model simply adds a linear layer with softmax activation to predict scores
|
@@ -29,27 +30,18 @@ def build_torch_ner_model(
|
|
29 |
nO (int or None): The number of tags to output. Inferred from the data if None.
|
30 |
RETURNS (Model[List[Doc], List[Floats2d]]): Initialized Model
|
31 |
"""
|
32 |
-
|
33 |
-
|
34 |
-
|
35 |
-
|
36 |
-
t2v_width = listener.maybe_get_dim("nO") if listener else None
|
37 |
-
#print(t2v_width, hidden_width, nO, dropout)
|
38 |
-
t2v_width = 768
|
39 |
-
#print(t2v_width, hidden_width, nO, dropout)
|
40 |
torch_model = TorchEntityRecognizer(t2v_width, hidden_width, nO, dropout)
|
41 |
-
#print("torch_model - ",torch_model)
|
42 |
wrapped_pt_model = PyTorchWrapper(torch_model)
|
43 |
-
#print("wrapped")
|
44 |
wrapped_pt_model.attrs["set_dropout_rate"] = torch_model.set_dropout_rate
|
45 |
-
#print("set dropout")
|
46 |
|
47 |
model = chain(tok2vec, with_array(wrapped_pt_model))
|
48 |
-
#print(model.param_names)
|
49 |
model.set_ref("tok2vec", tok2vec)
|
50 |
model.set_ref("torch_model", wrapped_pt_model)
|
51 |
model.init = init
|
52 |
-
#print("Completed build_torch_ner_model")
|
53 |
return model
|
54 |
|
55 |
|
@@ -65,39 +57,25 @@ def init(
|
|
65 |
RETURNS (Model[List[Doc], List[Floats2d]]): Initialized Model
|
66 |
"""
|
67 |
|
68 |
-
#print("Entered init - ")
|
69 |
tok2vec = model.get_ref("tok2vec")
|
70 |
-
|
71 |
torch_model = model.get_ref("torch_model")
|
72 |
-
#print(torch_model)
|
73 |
-
|
74 |
-
#print("Ref names - ",model.ref_names)
|
75 |
-
#print(tok2vec.dim_names,tok2vec.name)
|
76 |
-
#print(torch_model.dim_names,torch_model.name)
|
77 |
listener = tok2vec.maybe_get_ref("listener")
|
78 |
-
|
79 |
t2v_width = listener.maybe_get_dim("nO") if listener else None
|
80 |
-
|
81 |
if t2v_width:
|
82 |
-
#print(torch_model.shims[0]._model)
|
83 |
-
#print("Searching - ",torch_model.maybe_get_dim("nI"))
|
84 |
torch_model.shims[0]._model.set_input_shape(t2v_width)
|
85 |
torch_model.set_dim("nI", t2v_width)
|
86 |
-
#print(torch_model.dim_names)
|
87 |
|
88 |
if Y is not None:
|
89 |
nO = len(Y)
|
90 |
-
#print(nO)
|
91 |
torch_model.shims[0]._model.set_output_shape(nO)
|
92 |
torch_model.set_dim("nO", nO)
|
93 |
-
#print(torch_model)
|
94 |
|
95 |
tok2vec = model.get_ref("tok2vec")
|
96 |
tok2vec.initialize()
|
97 |
-
#print(tok2vec)
|
98 |
torch_model = model.get_ref("torch_model")
|
99 |
-
#print("Found - ",torch_model.get_dim("nI"))
|
100 |
-
#print("Exit")
|
101 |
return model
|
102 |
|
103 |
|
@@ -110,7 +88,6 @@ def is_dropout_module(
|
|
110 |
dropout_modules (List[nn.Module], optional): List of Modules that count as Dropout layers.
|
111 |
RETURNS (bool): True if module is a Dropout layer.
|
112 |
"""
|
113 |
-
#print("Entered is_dropout_module - ")
|
114 |
for m in dropout_modules:
|
115 |
if isinstance(module, m):
|
116 |
return True
|
@@ -130,7 +107,6 @@ class TorchEntityRecognizer(nn.Module):
|
|
130 |
super(TorchEntityRecognizer, self).__init__()
|
131 |
|
132 |
# Just for initialization of PyTorch layer. Output shape set during Model.init
|
133 |
-
#print("Entered TorchEntityRecognizer.__init__ - ")
|
134 |
nI = nI or 1
|
135 |
nO = nO or 1
|
136 |
|
@@ -147,14 +123,12 @@ class TorchEntityRecognizer(nn.Module):
|
|
147 |
}
|
148 |
)
|
149 |
)
|
150 |
-
#print(self.model)
|
151 |
|
152 |
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
|
153 |
"""Forward pass of the model.
|
154 |
inputs (torch.Tensor): Batch of outputs from spaCy tok2vec layer
|
155 |
RETURNS (torch.Tensor): Batch of results with a score for each tag for each token
|
156 |
"""
|
157 |
-
#print("Entered TorchEntityRecognizer.forward - ")
|
158 |
return self.model(inputs)
|
159 |
|
160 |
def _set_layer_shape(self, name: str, nI: int, nO: int):
|
@@ -163,39 +137,31 @@ class TorchEntityRecognizer(nn.Module):
|
|
163 |
nI (int): New input shape
|
164 |
nO (int): New output shape
|
165 |
"""
|
166 |
-
#print("Entered TorchEntityRecognizer._set_layer_shape - ",nO, nI)
|
167 |
with torch.no_grad():
|
168 |
layer = getattr(self.model, name)
|
169 |
#print(layer)
|
170 |
layer.out_features = nO
|
171 |
layer.weight = nn.Parameter(torch.Tensor(nO, nI))
|
172 |
-
#print(layer.weight.shape)
|
173 |
if layer.bias is not None:
|
174 |
layer.bias = nn.Parameter(torch.Tensor(nO))
|
175 |
-
#print(layer)
|
176 |
layer.reset_parameters()
|
177 |
-
#print(layer.weight.shape)
|
178 |
-
#print(layer)
|
179 |
|
180 |
def set_input_shape(self, nI: int):
|
181 |
"""Dynamically set the shape of the input layer
|
182 |
nI (int): New input layer shape
|
183 |
"""
|
184 |
-
#print("Entered TorchEntityRecognizer.set_input_shape - ",nI, self.nH)
|
185 |
self._set_layer_shape("input_layer", nI, self.nH)
|
186 |
|
187 |
def set_output_shape(self, nO: int):
|
188 |
"""Dynamically set the shape of the output layer
|
189 |
nO (int): New output layer shape
|
190 |
"""
|
191 |
-
#print("Entered TorchEntityRecognizer.set_output_shape - ", self.nH, nO)
|
192 |
self._set_layer_shape("output_layer", self.nH, nO)
|
193 |
|
194 |
def set_dropout_rate(self, dropout: float):
|
195 |
"""Set the dropout rate of all Dropout layers in the model.
|
196 |
dropout (float): Dropout rate to set
|
197 |
"""
|
198 |
-
#print("Entered TorchEntityRecognizer.set_dropout_rate - ")
|
199 |
dropout_layers = [
|
200 |
module for module in self.modules() if is_dropout_module(module)
|
201 |
]
|
|
|
21 |
hidden_width: int,
|
22 |
dropout: Optional[float] = None,
|
23 |
nO: Optional[int] = None,
|
24 |
+
width: Optional[int] = None,
|
25 |
) -> Model[List[Doc], List[Floats2d]]:
|
26 |
"""Build a tagger model, using a provided token-to-vector component. The tagger
|
27 |
model simply adds a linear layer with softmax activation to predict scores
|
|
|
30 |
nO (int or None): The number of tags to output. Inferred from the data if None.
|
31 |
RETURNS (Model[List[Doc], List[Floats2d]]): Initialized Model
|
32 |
"""
|
33 |
+
if width == None :
|
34 |
+
t2v_width = tok2vec.maybe_get_dim("nO")
|
35 |
+
else :
|
36 |
+
t2v_width = width
|
|
|
|
|
|
|
|
|
37 |
torch_model = TorchEntityRecognizer(t2v_width, hidden_width, nO, dropout)
|
|
|
38 |
wrapped_pt_model = PyTorchWrapper(torch_model)
|
|
|
39 |
wrapped_pt_model.attrs["set_dropout_rate"] = torch_model.set_dropout_rate
|
|
|
40 |
|
41 |
model = chain(tok2vec, with_array(wrapped_pt_model))
|
|
|
42 |
model.set_ref("tok2vec", tok2vec)
|
43 |
model.set_ref("torch_model", wrapped_pt_model)
|
44 |
model.init = init
|
|
|
45 |
return model
|
46 |
|
47 |
|
|
|
57 |
RETURNS (Model[List[Doc], List[Floats2d]]): Initialized Model
|
58 |
"""
|
59 |
|
|
|
60 |
tok2vec = model.get_ref("tok2vec")
|
61 |
+
|
62 |
torch_model = model.get_ref("torch_model")
|
|
|
|
|
|
|
|
|
|
|
63 |
listener = tok2vec.maybe_get_ref("listener")
|
64 |
+
|
65 |
t2v_width = listener.maybe_get_dim("nO") if listener else None
|
66 |
+
|
67 |
if t2v_width:
|
|
|
|
|
68 |
torch_model.shims[0]._model.set_input_shape(t2v_width)
|
69 |
torch_model.set_dim("nI", t2v_width)
|
|
|
70 |
|
71 |
if Y is not None:
|
72 |
nO = len(Y)
|
|
|
73 |
torch_model.shims[0]._model.set_output_shape(nO)
|
74 |
torch_model.set_dim("nO", nO)
|
|
|
75 |
|
76 |
tok2vec = model.get_ref("tok2vec")
|
77 |
tok2vec.initialize()
|
|
|
78 |
torch_model = model.get_ref("torch_model")
|
|
|
|
|
79 |
return model
|
80 |
|
81 |
|
|
|
88 |
dropout_modules (List[nn.Module], optional): List of Modules that count as Dropout layers.
|
89 |
RETURNS (bool): True if module is a Dropout layer.
|
90 |
"""
|
|
|
91 |
for m in dropout_modules:
|
92 |
if isinstance(module, m):
|
93 |
return True
|
|
|
107 |
super(TorchEntityRecognizer, self).__init__()
|
108 |
|
109 |
# Just for initialization of PyTorch layer. Output shape set during Model.init
|
|
|
110 |
nI = nI or 1
|
111 |
nO = nO or 1
|
112 |
|
|
|
123 |
}
|
124 |
)
|
125 |
)
|
|
|
126 |
|
127 |
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
|
128 |
"""Forward pass of the model.
|
129 |
inputs (torch.Tensor): Batch of outputs from spaCy tok2vec layer
|
130 |
RETURNS (torch.Tensor): Batch of results with a score for each tag for each token
|
131 |
"""
|
|
|
132 |
return self.model(inputs)
|
133 |
|
134 |
def _set_layer_shape(self, name: str, nI: int, nO: int):
|
|
|
137 |
nI (int): New input shape
|
138 |
nO (int): New output shape
|
139 |
"""
|
|
|
140 |
with torch.no_grad():
|
141 |
layer = getattr(self.model, name)
|
142 |
#print(layer)
|
143 |
layer.out_features = nO
|
144 |
layer.weight = nn.Parameter(torch.Tensor(nO, nI))
|
|
|
145 |
if layer.bias is not None:
|
146 |
layer.bias = nn.Parameter(torch.Tensor(nO))
|
|
|
147 |
layer.reset_parameters()
|
|
|
|
|
148 |
|
149 |
def set_input_shape(self, nI: int):
|
150 |
"""Dynamically set the shape of the input layer
|
151 |
nI (int): New input layer shape
|
152 |
"""
|
|
|
153 |
self._set_layer_shape("input_layer", nI, self.nH)
|
154 |
|
155 |
def set_output_shape(self, nO: int):
|
156 |
"""Dynamically set the shape of the output layer
|
157 |
nO (int): New output layer shape
|
158 |
"""
|
|
|
159 |
self._set_layer_shape("output_layer", self.nH, nO)
|
160 |
|
161 |
def set_dropout_rate(self, dropout: float):
|
162 |
"""Set the dropout rate of all Dropout layers in the model.
|
163 |
dropout (float): Dropout rate to set
|
164 |
"""
|
|
|
165 |
dropout_layers = [
|
166 |
module for module in self.modules() if is_dropout_module(module)
|
167 |
]
|