Update scripts/torch_ner_model.py

scripts/torch_ner_model.py

@@ -21,6 +21,7 @@ def build_torch_ner_model(
     hidden_width: int,
     dropout: Optional[float] = None,
     nO: Optional[int] = None,
+    width: Optional[int] = None,
 ) -> Model[List[Doc], List[Floats2d]]:
     """Build a tagger model, using a provided token-to-vector component. The tagger
     model simply adds a linear layer with softmax activation to predict scores
@@ -29,27 +30,18 @@ def build_torch_ner_model(
     nO (int or None): The number of tags to output. Inferred from the data if None.
     RETURNS (Model[List[Doc], List[Floats2d]]): Initialized Model
     """
-    ##print("Entered build_torch_ner_model - ")
-    #print(tok2vec.dim_names,tok2vec.name)
-    listener = tok2vec.maybe_get_ref("listener")
-    #print(listener.maybe_get_dim("nI"))
-    t2v_width = listener.maybe_get_dim("nO") if listener else None
-    #print(t2v_width, hidden_width, nO, dropout)
-    t2v_width = 768
-    #print(t2v_width, hidden_width, nO, dropout)
+    if width == None :
+        t2v_width = tok2vec.maybe_get_dim("nO")
+    else :
+        t2v_width = width
     torch_model = TorchEntityRecognizer(t2v_width, hidden_width, nO, dropout)
-    #print("torch_model - ",torch_model)
     wrapped_pt_model = PyTorchWrapper(torch_model)
-    #print("wrapped")
     wrapped_pt_model.attrs["set_dropout_rate"] = torch_model.set_dropout_rate
-    #print("set dropout")
     
     model = chain(tok2vec, with_array(wrapped_pt_model))
-    #print(model.param_names)
     model.set_ref("tok2vec", tok2vec)
     model.set_ref("torch_model", wrapped_pt_model)
     model.init = init
-    #print("Completed build_torch_ner_model")
     return model
 
 
@@ -65,39 +57,25 @@ def init(
     RETURNS (Model[List[Doc], List[Floats2d]]): Initialized Model
     """
 
-    #print("Entered init - ")
     tok2vec = model.get_ref("tok2vec")
-    #print(tok2vec.ref_names)
+
     torch_model = model.get_ref("torch_model")
-    #print(torch_model)
-    
-    #print("Ref names - ",model.ref_names)
-    #print(tok2vec.dim_names,tok2vec.name)
-    #print(torch_model.dim_names,torch_model.name)
     listener = tok2vec.maybe_get_ref("listener")
-    #print(listener)
+
     t2v_width = listener.maybe_get_dim("nO") if listener else None
-    #print(t2v_width," - ",Y)
+
     if t2v_width:
-        #print(torch_model.shims[0]._model)
-        #print("Searching - ",torch_model.maybe_get_dim("nI"))
         torch_model.shims[0]._model.set_input_shape(t2v_width)
         torch_model.set_dim("nI", t2v_width)
-        #print(torch_model.dim_names)
 
     if Y is not None:
         nO = len(Y)
-        #print(nO)
         torch_model.shims[0]._model.set_output_shape(nO)
         torch_model.set_dim("nO", nO)
-        #print(torch_model)
 
     tok2vec = model.get_ref("tok2vec")
     tok2vec.initialize()
-    #print(tok2vec)
     torch_model = model.get_ref("torch_model")
-    #print("Found - ",torch_model.get_dim("nI"))
-    #print("Exit")
     return model
 
 
@@ -110,7 +88,6 @@ def is_dropout_module(
     dropout_modules (List[nn.Module], optional): List of Modules that count as Dropout layers.
     RETURNS (bool): True if module is a Dropout layer.
     """
-    #print("Entered is_dropout_module - ")
     for m in dropout_modules:
         if isinstance(module, m):
             return True
@@ -130,7 +107,6 @@ class TorchEntityRecognizer(nn.Module):
         super(TorchEntityRecognizer, self).__init__()
 
         # Just for initialization of PyTorch layer. Output shape set during Model.init
-        #print("Entered TorchEntityRecognizer.__init__  - ")
         nI = nI or 1
         nO = nO or 1
 
@@ -147,14 +123,12 @@ class TorchEntityRecognizer(nn.Module):
                 }
             )
         )
-        #print(self.model)
 
     def forward(self, inputs: torch.Tensor) -> torch.Tensor:
         """Forward pass of the model.
         inputs (torch.Tensor): Batch of outputs from spaCy tok2vec layer
         RETURNS (torch.Tensor): Batch of results with a score for each tag for each token
         """
-        #print("Entered TorchEntityRecognizer.forward  - ")
         return self.model(inputs)
 
     def _set_layer_shape(self, name: str, nI: int, nO: int):
@@ -163,39 +137,31 @@ class TorchEntityRecognizer(nn.Module):
         nI (int): New input shape
         nO (int): New output shape
         """
-        #print("Entered TorchEntityRecognizer._set_layer_shape  - ",nO, nI)
         with torch.no_grad():
             layer = getattr(self.model, name)
             #print(layer)
             layer.out_features = nO
             layer.weight = nn.Parameter(torch.Tensor(nO, nI))
-            #print(layer.weight.shape)
             if layer.bias is not None:
                 layer.bias = nn.Parameter(torch.Tensor(nO))
-            #print(layer)
             layer.reset_parameters()
-            #print(layer.weight.shape)
-            #print(layer)
 
     def set_input_shape(self, nI: int):
         """Dynamically set the shape of the input layer
         nI (int): New input layer shape
         """
-        #print("Entered TorchEntityRecognizer.set_input_shape  - ",nI, self.nH)
         self._set_layer_shape("input_layer", nI, self.nH)
 
     def set_output_shape(self, nO: int):
         """Dynamically set the shape of the output layer
         nO (int): New output layer shape
         """
-        #print("Entered TorchEntityRecognizer.set_output_shape  - ", self.nH, nO)
         self._set_layer_shape("output_layer", self.nH, nO)
 
     def set_dropout_rate(self, dropout: float):
         """Set the dropout rate of all Dropout layers in the model.
         dropout (float): Dropout rate to set
         """
-        #print("Entered TorchEntityRecognizer.set_dropout_rate  - ")
         dropout_layers = [
             module for module in self.modules() if is_dropout_module(module)
         ]