Refactor for easier testing

Author: sergey-misuk-valor

src/country_workspace/contrib/name_parser/parser.py

@@ -36,32 +36,50 @@ def forward(self, input_: torch.Tensor) -> torch.Tensor:
         return self.softmax(out)
 
 
-def get_parser(country_code: str) -> Parser:
+Alphabet = tuple[str, ...]
+ModelArgs = tuple[int, ...]
+UNKNOWN_CHAR = "_"
+
+
+def read_config(country_code: str) -> tuple[Alphabet, int, ModelArgs]:
     with (BASE_PATH / f"data/name_parser/models/{country_code}.txt").open() as f:
         lines = tuple(line.rstrip("\n") for line in f.readlines())
 
-    unknown = "_"
-    alphabet = tuple(lines[0])
-    alphabet_len = len(alphabet)
-    name_max_len = int(lines[1])
-    rnn_args = map(int, lines[2].split())
+    return (
+        tuple(lines[0]),
+        int(lines[1]),
+        tuple(map(int, lines[2].split())),
+    )
 
-    rnn = LSTM(*rnn_args, num_layers=2)
+
+def load_model(country_code: str, *args: int) -> nn.Module:
+    rnn = LSTM(*args, num_layers=2)
     rnn.load_state_dict(torch.load(BASE_PATH / f"data/name_parser/models/{country_code}.pt"))
     rnn.to(DEVICE)
     rnn.eval()
+    return rnn
 
-    def letter_to_index(letter: str) -> int:
-        return alphabet.index(letter) if letter in alphabet else alphabet.index(unknown)
 
-    oob = alphabet_len + 1
+def get_line_to_tensor_converter(alphabet: Alphabet, max_name_len: int) -> Callable[[str], torch.Tensor]:
+    oob = len(alphabet) + 1
+
+    def letter_to_index(letter: str) -> int:
+        return alphabet.index(letter) if letter in alphabet else alphabet.index(UNKNOWN_CHAR)
 
     def line_to_tensor(line: str) -> torch.Tensor:
-        tensor = torch.ones(name_max_len, dtype=torch.long) * oob
+        tensor = torch.ones(max_name_len, dtype=torch.long) * oob
         for li, letter in enumerate(line):
             tensor[li] = letter_to_index(letter)
         return tensor
 
+    return line_to_tensor
+
+
+def get_parser(country_code: str) -> Parser:
+    alphabet, max_name_len, rnn_args = read_config(country_code)
+    rnn = load_model(country_code, *rnn_args)
+    line_to_tensor = get_line_to_tensor_converter(alphabet, max_name_len)
+
     def parser(name: str) -> list[str]:
         name_tokens = [line_to_tensor(i) for i in name.split()]
         out = [rnn(i.unsqueeze(0).to(DEVICE)) for i in name_tokens]