Update handler.py

handler.py

@@ -12,6 +12,7 @@ from model import build_transformer
 
 warnings.simplefilter("ignore", category=FutureWarning)
 
+
 class EndpointHandler:
     def __init__(self, path: str = ""):
         """
@@ -48,51 +49,50 @@ class EndpointHandler:
         self.model.load_state_dict(checkpoint["model_state_dict"])
         self.model.eval()
 
-   def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
-    """
-    Process the incoming request and return the translation.
-    """
-    try:
-        inputs = data.get("inputs", "")
-        if not inputs:
-            return [{"error": "No 'inputs' provided in request"}]
-
-        source = self.tokenizer_src.encode(inputs)
-        source = torch.cat([
-            torch.tensor([self.tokenizer_src.token_to_id("[SOS]")], dtype=torch.int64),
-            torch.tensor(source.ids, dtype=torch.int64),
-            torch.tensor([self.tokenizer_src.token_to_id("[EOS]")], dtype=torch.int64),
-            torch.tensor([self.tokenizer_src.token_to_id("[PAD]")] * (350 - len(source.ids) - 2), dtype=torch.int64)
-        ], dim=0).to(self.device)
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+        Process the incoming request and return the translation.
+        """
+        try:
+            inputs = data.get("inputs", "")
+            if not inputs:
+                return [{"error": "No 'inputs' provided in request"}]
 
-        source_mask = (source != self.tokenizer_src.token_to_id("[PAD]")).unsqueeze(0).unsqueeze(1).int().to(self.device)
-        encoder_output = self.model.encode(source, source_mask)
+            source = self.tokenizer_src.encode(inputs)
+            source = torch.cat([
+                torch.tensor([self.tokenizer_src.token_to_id("[SOS]")], dtype=torch.int64),
+                torch.tensor(source.ids, dtype=torch.int64),
+                torch.tensor([self.tokenizer_src.token_to_id("[EOS]")], dtype=torch.int64),
+                torch.tensor([self.tokenizer_src.token_to_id("[PAD]")] * (350 - len(source.ids) - 2), dtype=torch.int64)
+            ], dim=0).to(self.device)
 
-        decoder_input = torch.empty(1, 1).fill_(self.tokenizer_tgt.token_to_id("[SOS]")).type_as(source).to(self.device)
-        predicted_words = []
+            source_mask = (source != self.tokenizer_src.token_to_id("[PAD]")).unsqueeze(0).unsqueeze(1).int().to(self.device)
+            encoder_output = self.model.encode(source, source_mask)
 
-        while decoder_input.size(1) < 350:
-            decoder_mask = torch.triu(
-                torch.ones((1, decoder_input.size(1), decoder_input.size(1))),
-                diagonal=1
-            ).type(torch.int).type_as(source_mask).to(self.device)
+            decoder_input = torch.empty(1, 1).fill_(self.tokenizer_tgt.token_to_id("[SOS]")).type_as(source).to(self.device)
+            predicted_words = []
 
-            out = self.model.decode(encoder_output, source_mask, decoder_input, decoder_mask)
-            prob = self.model.project(out[:, -1])
-            _, next_word = torch.max(prob, dim=1)
+            while decoder_input.size(1) < 350:
+                decoder_mask = torch.triu(
+                    torch.ones((1, decoder_input.size(1), decoder_input.size(1))),
+                    diagonal=1
+                ).type(torch.int).type_as(source_mask).to(self.device)
 
-            decoder_input = torch.cat(
-                [decoder_input, torch.empty(1, 1).type_as(source).fill_(next_word.item()).to(self.device)], dim=1)
+                out = self.model.decode(encoder_output, source_mask, decoder_input, decoder_mask)
+                prob = self.model.project(out[:, -1])
+                _, next_word = torch.max(prob, dim=1)
 
-            decoded_word = self.tokenizer_tgt.decode([next_word.item()])
-            if next_word == self.tokenizer_tgt.token_to_id("[EOS]"):
-                break
+                decoder_input = torch.cat(
+                    [decoder_input, torch.empty(1, 1).type_as(source).fill_(next_word.item()).to(self.device)], dim=1)
 
-            predicted_words.append(decoded_word)
+                decoded_word = self.tokenizer_tgt.decode([next_word.item()])
+                if next_word == self.tokenizer_tgt.token_to_id("[EOS]"):
+                    break
 
-        predicted_translation = " ".join(predicted_words).replace("[EOS]", "").strip()
+                predicted_words.append(decoded_word)
 
-        return [{"translation": predicted_translation}]
-    except Exception as e:
-        return [{"error": str(e)}]
+            predicted_translation = " ".join(predicted_words).replace("[EOS]", "").strip()
 
+            return [{"translation": predicted_translation}]
+        except Exception as e:
+            return [{"error": str(e)}]