fixed python module to check the id range.

google · Jan 10, 2021 · 0e6dfbf · 0e6dfbf
1 parent 3589bfb
commit 0e6dfbf
Show file tree

Hide file tree

Showing 3 changed files with 177 additions and 131 deletions.
diff --git a/python/src/sentencepiece/__init__.py b/python/src/sentencepiece/__init__.py
@@ -116,9 +116,6 @@ def SampleEncodeAsIds(self, input, nbest_size, alpha):
     def DecodePieces(self, pieces):
         return _sentencepiece.SentencePieceProcessor_DecodePieces(self, pieces)
 
-    def DecodeIds(self, ids):
-        return _sentencepiece.SentencePieceProcessor_DecodeIds(self, ids)
-
     def EncodeAsSerializedProto(self, input):
         return _sentencepiece.SentencePieceProcessor_EncodeAsSerializedProto(self, input)
 
@@ -131,9 +128,6 @@ def NBestEncodeAsSerializedProto(self, input, nbest_size):
     def DecodePiecesAsSerializedProto(self, pieces):
         return _sentencepiece.SentencePieceProcessor_DecodePiecesAsSerializedProto(self, pieces)
 
-    def DecodeIdsAsSerializedProto(self, ids):
-        return _sentencepiece.SentencePieceProcessor_DecodeIdsAsSerializedProto(self, ids)
-
     def GetPieceSize(self):
         return _sentencepiece.SentencePieceProcessor_GetPieceSize(self)
 
@@ -176,6 +170,12 @@ def serialized_model_proto(self):
     def LoadFromFile(self, arg):
         return _sentencepiece.SentencePieceProcessor_LoadFromFile(self, arg)
 
+    def DecodeIdsWithCheck(self, ids):
+        return _sentencepiece.SentencePieceProcessor_DecodeIdsWithCheck(self, ids)
+
+    def DecodeIdsAsSerializedProtoWithCheck(self, ids):
+        return _sentencepiece.SentencePieceProcessor_DecodeIdsAsSerializedProtoWithCheck(self, ids)
+
     def Init(self,
              model_file=None,
              model_proto=None,
@@ -310,15 +310,15 @@ def Decode(self, input):
       if not input:
         return self.DecodeIds([])
       elif type(input) is int:
-        return self.DecodeIds([input])
+        return self.DecodeIdsWithCheck([input])
       elif type(input) is str:
         return self.DecodePieces([input])
 
       def _decode(input):
         if not input:
           return self.DecodeIds([])
         if type(input[0]) is int:
-          return self.DecodeIds(input)
+          return self.DecodeIdsWithCheck(input)
         return self.DecodePieces(input)
 
       if type(input[0]) is list:
@@ -508,6 +508,8 @@ def _batched_func(self, arg):
 
 SentencePieceProcessor.Tokenize = SentencePieceProcessor.Encode
 SentencePieceProcessor.Detokenize = SentencePieceProcessor.Decode
+SentencePieceProcessor.DecodeIds = SentencePieceProcessor.DecodeIdsWithCheck
+SentencePieceProcessor.DecodeIdsAsSerializedProto = SentencePieceProcessor.DecodeIdsAsSerializedProtoWithCheck
 
 for m in [
     'PieceToId', 'IdToPiece', 'GetScore', 'IsUnknown', 'IsControl', 'IsUnused',

diff --git a/python/src/sentencepiece/sentencepiece.i b/python/src/sentencepiece/sentencepiece.i
@@ -176,6 +176,8 @@ class PySentenceIterator : public sentencepiece::SentenceIterator {
 %ignore sentencepiece::SentencePieceProcessor::SampleEncode;
 %ignore sentencepiece::SentencePieceProcessor::NBestEncode;
 %ignore sentencepiece::SentencePieceProcessor::Decode;
+%ignore sentencepiece::SentencePieceProcessor::DecodeIds;
+%ignore sentencepiece::SentencePieceProcessor::DecodeIdsAsSerializedProto;
 %ignore sentencepiece::SentencePieceProcessor::model_proto;
 %ignore sentencepiece::SentencePieceProcessor::Load;
 %ignore sentencepiece::SentencePieceProcessor::LoadOrDie;
@@ -196,6 +198,28 @@ class PySentenceIterator : public sentencepiece::SentenceIterator {
     return $self->Load(arg);
   }
 
+  std::string DecodeIdsWithCheck(
+      const std::vector<int> &ids) const {
+    const int num_pieces = $self->GetPieceSize(); 
+    for (int id : ids)
+      if (id < 0 || id >= num_pieces)
+        throw sentencepiece::util::Status(
+            sentencepiece::util::StatusCode::kOutOfRange,
+            "piece id is out of range.");
+    return $self->DecodeIds(ids);
+  }
+
+  util::bytes DecodeIdsAsSerializedProtoWithCheck(
+      const std::vector<int> &ids) const {
+    const int num_pieces = $self->GetPieceSize(); 
+    for (int id : ids)
+      if (id < 0 || id >= num_pieces)
+        throw sentencepiece::util::Status(
+            sentencepiece::util::StatusCode::kOutOfRange,
+            "piece id is out of range.");
+    return $self->DecodeIdsAsSerializedProto(ids);
+  }
+
 %pythoncode {
   def Init(self,
            model_file=None,
@@ -331,15 +355,15 @@ class PySentenceIterator : public sentencepiece::SentenceIterator {
     if not input:
       return self.DecodeIds([])
     elif type(input) is int:
-      return self.DecodeIds([input])
+      return self.DecodeIdsWithCheck([input])
     elif type(input) is str:
       return self.DecodePieces([input])
 
     def _decode(input):
       if not input:
         return self.DecodeIds([])
       if type(input[0]) is int:
-        return self.DecodeIds(input)
+        return self.DecodeIdsWithCheck(input)
       return self.DecodePieces(input)
 
     if type(input[0]) is list:
@@ -707,6 +731,8 @@ setattr(SentencePieceProcessor, '__init__', SentencePieceProcessor.Init)
 
 SentencePieceProcessor.Tokenize = SentencePieceProcessor.Encode
 SentencePieceProcessor.Detokenize = SentencePieceProcessor.Decode
+SentencePieceProcessor.DecodeIds = SentencePieceProcessor.DecodeIdsWithCheck
+SentencePieceProcessor.DecodeIdsAsSerializedProto = SentencePieceProcessor.DecodeIdsAsSerializedProtoWithCheck
 
 for m in [
     'PieceToId', 'IdToPiece', 'GetScore', 'IsUnknown', 'IsControl', 'IsUnused',