sentencepiece 0.0.1 → 0.1.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 5fea5ff13a05fa657877f461944ccb7f5f4b8b6568d25c7581f795ad19514d9a
-  data.tar.gz: ecfb8603781a0b9129ea68d94272739c558d07a1aaff8f35b5b76c1bcbb4ee87
+  metadata.gz: 63aaf69a72781f0087072db8197b9ae3a126b63e632b52bded5b483ff7b2c6b4
+  data.tar.gz: 5ffbd48ed587ef57fd3b427615cc21f4d79804ea90598f9db2178ceff702ca57
 SHA512:
-  metadata.gz: fe928e1d7810a77f34ccce58b4feea2754ba533fbe1a3e4a0f8ccfbbfb0bdb5b0366550ce5844e6df50ad68a261fe8e9a6c047dde415683bb072d91c39f91868
-  data.tar.gz: 5247986788a8a7c6678b65d23d235647518e25f32e1d888246255d556b20e82668331de0b7a8d71998be06163b1cded08be51af7a2b23719499f2aac79e6cd72
+  metadata.gz: 61b91c838bbdb2f7c47166036b890b545060361e91a66e6cf416206f53d84595cc74052a38df083ad507a6f219d0b2b8ba3667095d2d3514b5ba0c3660967ee5
+  data.tar.gz: 467c2f252e23fbfe684486d42af9fe0eedd058dd99361fc5652bb83412ae801571e5e25879b2f1bb6734c997b9c04f4287378c595afa3698b06885631c8fdcae

data/CHANGELOG.md

@@ -1,5 +1,15 @@
 ## [Unreleased]
 
+## [0.1.0] - 2023-03-26
+
+- Add API documentation.
+- Add type signatures.
+
+## [0.0.2] - 2023-03-26
+
+- Add SentencePieceTrainer class.
+- Add some encoding and decoding methods to SentencePieceProcessor.
+
 ## [0.0.1] - 2023-03-21
 
 - Initial release

data/README.md

@@ -1,12 +1,13 @@
 # sentencepiece.rb
 
+[![Build Status](https://github.com/yoshoku/sentencepiece.rb/actions/workflows/main.yml/badge.svg)](https://github.com/yoshoku/sentencepiece.rb/actions/workflows/main.yml)
+[![Gem Version](https://badge.fury.io/rb/sentencepiece.svg)](https://badge.fury.io/rb/sentencepiece)
 [![License](https://img.shields.io/badge/License-Apache%202.0-yellowgreen.svg)](https://github.com/yoshoku/sentencepiece.rb/blob/main/LICENSE.txt)
+[![Documentation](https://img.shields.io/badge/api-reference-blue.svg)](https://yoshoku.github.io/sentencepiece.rb/doc/)
 
 sentencepiece.rb provides Ruby bindings for the [SentencePiece](https://github.com/google/sentencepiece),
 an unsupervised text tokenizer and detokenizer for neural network-based text generation.
 
-It is still under development and may undergo many changes in the future.
-
 ## Installation
 
 Install SentencePiece using your OS package manager;

data/ext/sentencepiece/sentencepiece.cpp

@@ -22,4 +22,5 @@ extern "C" void Init_sentencepiece(void) {
   rb_mSentencePiece = rb_define_module("SentencePiece");
   rb_eSentencePieceError = rb_define_class_under(rb_mSentencePiece, "Error", rb_eRuntimeError);
   RbSentencePieceProcessor::define_class(rb_mSentencePiece);
+  RbSentencePieceTrainer::define_class(rb_mSentencePiece);
 }

data/ext/sentencepiece/sentencepiece.hpp

@@ -27,6 +27,7 @@
 VALUE rb_mSentencePiece;
 VALUE rb_eSentencePieceError;
 VALUE rb_cSentencePieceProcessor;
+VALUE rb_cSentencePieceTrainer;
 
 class RbSentencePieceProcessor {
 public:
@@ -57,10 +58,27 @@ public:
     rb_define_method(rb_cSentencePieceProcessor, "initialize", RUBY_METHOD_FUNC(_sentencepiece_processor_init), -1);
     rb_define_method(rb_cSentencePieceProcessor, "load", RUBY_METHOD_FUNC(_sentencepiece_processor_load), 1);
     rb_define_method(rb_cSentencePieceProcessor, "encode", RUBY_METHOD_FUNC(_sentencepiece_processor_encode), -1);
+    rb_define_method(rb_cSentencePieceProcessor, "encode_as_ids", RUBY_METHOD_FUNC(_sentencepiece_processor_encode_as_ids), 1);
+    rb_define_method(rb_cSentencePieceProcessor, "encode_as_pieces", RUBY_METHOD_FUNC(_sentencepiece_processor_encode_as_pieces), 1);
+    rb_define_method(rb_cSentencePieceProcessor, "nbest_encode_as_pieces", RUBY_METHOD_FUNC(_sentencepiece_processor_nbest_encode_as_pieces), -1);
+    rb_define_method(rb_cSentencePieceProcessor, "nbest_encode_as_ids", RUBY_METHOD_FUNC(_sentencepiece_processor_nbest_encode_as_ids), -1);
+    rb_define_method(rb_cSentencePieceProcessor, "sample_encode_as_pieces", RUBY_METHOD_FUNC(_sentencepiece_processor_sample_encode_as_pieces), -1);
+    rb_define_method(rb_cSentencePieceProcessor, "sample_encode_as_ids", RUBY_METHOD_FUNC(_sentencepiece_processor_sample_encode_as_ids), -1);
     rb_define_method(rb_cSentencePieceProcessor, "decode", RUBY_METHOD_FUNC(_sentencepiece_processor_decode), -1);
+    rb_define_method(rb_cSentencePieceProcessor, "decode_pieces", RUBY_METHOD_FUNC(_sentencepiece_processor_decode_pieces), 1);
+    rb_define_method(rb_cSentencePieceProcessor, "decode_ids", RUBY_METHOD_FUNC(_sentencepiece_processor_decode_ids), 1);
+    rb_define_method(rb_cSentencePieceProcessor, "encode_as_serialized_proto", RUBY_METHOD_FUNC(_sentencepiece_processor_encode_as_serialized_proto), 1);
+    rb_define_method(rb_cSentencePieceProcessor, "sample_encode_as_serialized_proto", RUBY_METHOD_FUNC(_sentencepiece_processor_sample_encode_as_serialized_proto), -1);
+    rb_define_method(rb_cSentencePieceProcessor, "nbest_encode_as_serialized_proto", RUBY_METHOD_FUNC(_sentencepiece_processor_nbest_encode_as_serialized_proto), -1);
+    rb_define_method(rb_cSentencePieceProcessor, "decode_pieces_as_serialized_proto", RUBY_METHOD_FUNC(_sentencepiece_processor_decode_pieces_as_serialized_proto), 1);
+    rb_define_method(rb_cSentencePieceProcessor, "decode_ids_as_serialized_proto", RUBY_METHOD_FUNC(_sentencepiece_processor_decode_ids_as_serialized_proto), 1);
     rb_define_method(rb_cSentencePieceProcessor, "piece_size", RUBY_METHOD_FUNC(_sentencepiece_processor_piece_size), 0);
     rb_define_method(rb_cSentencePieceProcessor, "piece_to_id", RUBY_METHOD_FUNC(_sentencepiece_processor_piece_to_id), 1);
     rb_define_method(rb_cSentencePieceProcessor, "id_to_piece", RUBY_METHOD_FUNC(_sentencepiece_processor_id_to_piece), 1);
+    rb_define_method(rb_cSentencePieceProcessor, "unk_id", RUBY_METHOD_FUNC(_sentencepiece_processor_unk_id), 0);
+    rb_define_method(rb_cSentencePieceProcessor, "bos_id", RUBY_METHOD_FUNC(_sentencepiece_processor_bos_id), 0);
+    rb_define_method(rb_cSentencePieceProcessor, "eos_id", RUBY_METHOD_FUNC(_sentencepiece_processor_eos_id), 0);
+    rb_define_method(rb_cSentencePieceProcessor, "pad_id", RUBY_METHOD_FUNC(_sentencepiece_processor_pad_id), 0);
     return rb_cSentencePieceProcessor;
   };
 
@@ -193,6 +211,182 @@ private:
     return output;
   };
 
+  static VALUE _sentencepiece_processor_encode_as_ids(VALUE self, VALUE text) {
+    if (!RB_TYPE_P(text, T_STRING)) {
+      rb_raise(rb_eArgError, "expected text to be a String");
+      return Qnil;
+    }
+
+    sentencepiece::SentencePieceProcessor* ptr = get_sentencepiece_processor(self);
+    const std::vector<int> ids = ptr->EncodeAsIds(StringValueCStr(text));
+    VALUE output = rb_ary_new();
+    for (const int idx : ids) {
+      rb_ary_push(output, INT2NUM(idx));
+    }
+
+    RB_GC_GUARD(text);
+    return output;
+  };
+
+  static VALUE _sentencepiece_processor_encode_as_pieces(VALUE self, VALUE text) {
+    if (!RB_TYPE_P(text, T_STRING)) {
+      rb_raise(rb_eArgError, "expected text to be a String");
+      return Qnil;
+    }
+
+    sentencepiece::SentencePieceProcessor* ptr = get_sentencepiece_processor(self);
+    const std::vector<std::string> pieces = ptr->EncodeAsPieces(StringValueCStr(text));
+    VALUE output = rb_ary_new();
+    for (const std::string& token : pieces) {
+      rb_ary_push(output, rb_utf8_str_new_cstr(token.c_str()));
+    }
+
+    RB_GC_GUARD(text);
+    return output;
+  };
+
+  static VALUE _sentencepiece_processor_nbest_encode_as_pieces(int argc, VALUE* argv, VALUE self) {
+    VALUE kw_args = Qnil;
+    ID kw_table[1] = { rb_intern("nbest_size") };
+    VALUE kw_values[1] = { Qundef };
+
+    VALUE text = Qnil;
+    rb_scan_args(argc, argv, "1:", &text, &kw_args);
+    rb_get_kwargs(kw_args, kw_table, 1, 0, kw_values);
+
+    if (!RB_TYPE_P(text, T_STRING)) {
+      rb_raise(rb_eArgError, "expected text to be a String");
+      return Qnil;
+    }
+    if (!RB_INTEGER_TYPE_P(kw_values[0])) {
+      rb_raise(rb_eArgError, "expected nbest_size to be an Integer");
+      return Qnil;
+    }
+
+    const int nbest_size = NUM2INT(kw_values[0]);
+    sentencepiece::SentencePieceProcessor* ptr = get_sentencepiece_processor(self);
+    const std::vector<std::vector<std::string>> pieces = ptr->NBestEncodeAsPieces(StringValueCStr(text), nbest_size);
+
+    VALUE output = rb_ary_new();
+    for (const std::vector<std::string>& tokens : pieces) {
+      VALUE sub_output = rb_ary_new();
+      for (const std::string& token : tokens) {
+        rb_ary_push(sub_output, rb_utf8_str_new_cstr(token.c_str()));
+      }
+      rb_ary_push(output, sub_output);
+    }
+
+    RB_GC_GUARD(text);
+    return output;
+  };
+
+  static VALUE _sentencepiece_processor_nbest_encode_as_ids(int argc, VALUE* argv, VALUE self) {
+    VALUE kw_args = Qnil;
+    ID kw_table[1] = { rb_intern("nbest_size") };
+    VALUE kw_values[1] = { Qundef };
+
+    VALUE text = Qnil;
+    rb_scan_args(argc, argv, "1:", &text, &kw_args);
+    rb_get_kwargs(kw_args, kw_table, 1, 0, kw_values);
+
+    if (!RB_TYPE_P(text, T_STRING)) {
+      rb_raise(rb_eArgError, "expected text to be a String");
+      return Qnil;
+    }
+    if (!RB_INTEGER_TYPE_P(kw_values[0])) {
+      rb_raise(rb_eArgError, "expected nbest_size to be an Integer");
+      return Qnil;
+    }
+
+    const int nbest_size = NUM2INT(kw_values[0]);
+    sentencepiece::SentencePieceProcessor* ptr = get_sentencepiece_processor(self);
+    const std::vector<std::vector<int>> id_set = ptr->NBestEncodeAsIds(StringValueCStr(text), nbest_size);
+
+    VALUE output = rb_ary_new();
+    for (const std::vector<int> ids : id_set) {
+      VALUE sub_output = rb_ary_new();
+      for (const int idx : ids) {
+        rb_ary_push(sub_output, INT2NUM(idx));
+      }
+      rb_ary_push(output, sub_output);
+    }
+
+    RB_GC_GUARD(text);
+    return output;
+  };
+
+  static VALUE _sentencepiece_processor_sample_encode_as_pieces(int argc, VALUE* argv, VALUE self) {
+    VALUE kw_args = Qnil;
+    ID kw_table[2] = { rb_intern("nbest_size"), rb_intern("alpha") };
+    VALUE kw_values[2] = { Qundef, Qundef };
+
+    VALUE text = Qnil;
+    rb_scan_args(argc, argv, "1:", &text, &kw_args);
+    rb_get_kwargs(kw_args, kw_table, 2, 0, kw_values);
+
+    if (!RB_TYPE_P(text, T_STRING)) {
+      rb_raise(rb_eArgError, "expected text to be a String");
+      return Qnil;
+    }
+    if (!RB_INTEGER_TYPE_P(kw_values[0])) {
+      rb_raise(rb_eArgError, "expected nbest_size to be an Integer");
+      return Qnil;
+    }
+    if (!RB_INTEGER_TYPE_P(kw_values[1]) && !RB_FLOAT_TYPE_P(kw_values[1])) {
+      rb_raise(rb_eArgError, "expected alpha to be a Float");
+      return Qnil;
+    }
+
+    const int nbest_size = NUM2INT(kw_values[0]);
+    const float alpha = NUM2DBL(kw_values[1]);
+    sentencepiece::SentencePieceProcessor* ptr = get_sentencepiece_processor(self);
+    const std::vector<std::string> pieces = ptr->SampleEncodeAsPieces(StringValueCStr(text), nbest_size, alpha);
+
+    VALUE output = rb_ary_new();
+    for (const std::string& token : pieces) {
+      rb_ary_push(output, rb_utf8_str_new_cstr(token.c_str()));
+    }
+
+    RB_GC_GUARD(text);
+    return output;
+  };
+
+  static VALUE _sentencepiece_processor_sample_encode_as_ids(int argc, VALUE* argv, VALUE self) {
+    VALUE kw_args = Qnil;
+    ID kw_table[2] = { rb_intern("nbest_size"), rb_intern("alpha") };
+    VALUE kw_values[2] = { Qundef, Qundef };
+
+    VALUE text = Qnil;
+    rb_scan_args(argc, argv, "1:", &text, &kw_args);
+    rb_get_kwargs(kw_args, kw_table, 2, 0, kw_values);
+
+    if (!RB_TYPE_P(text, T_STRING)) {
+      rb_raise(rb_eArgError, "expected text to be a String");
+      return Qnil;
+    }
+    if (!RB_INTEGER_TYPE_P(kw_values[0])) {
+      rb_raise(rb_eArgError, "expected nbest_size to be an Integer");
+      return Qnil;
+    }
+    if (!RB_INTEGER_TYPE_P(kw_values[1]) && !RB_FLOAT_TYPE_P(kw_values[1])) {
+      rb_raise(rb_eArgError, "expected alpha to be a Float");
+      return Qnil;
+    }
+
+    const int nbest_size = NUM2INT(kw_values[0]);
+    const float alpha = NUM2DBL(kw_values[1]);
+    sentencepiece::SentencePieceProcessor* ptr = get_sentencepiece_processor(self);
+    const std::vector<int> ids = ptr->SampleEncodeAsIds(StringValueCStr(text), nbest_size, alpha);
+
+    VALUE output = rb_ary_new();
+    for (const int idx : ids) {
+      rb_ary_push(output, INT2NUM(idx));
+    }
+
+    RB_GC_GUARD(text);
+    return output;
+  };
+
   static VALUE _sentencepiece_processor_decode(int argc, VALUE* argv, VALUE self) {
     VALUE kw_args = Qnil;
     ID kw_table[1] = { rb_intern("out_type") };
@@ -204,7 +398,7 @@ private:
     VALUE out_type = kw_values[0] != Qundef ? kw_values[0] : rb_str_new_cstr("int");
 
     if (!RB_TYPE_P(pieces, T_ARRAY)) {
-      rb_raise(rb_eArgError, "expected out_type to be an Array");
+      rb_raise(rb_eArgError, "expected pieces to be an Array");
       return Qnil;
     }
     if (strcmp(StringValueCStr(out_type), "str") != 0 && strcmp(StringValueCStr(out_type), "int") != 0) {
@@ -232,7 +426,7 @@ private:
           rb_raise(rb_eSentencePieceError, "%s", status.message());
           return Qfalse;
         }
-        output = rb_str_new_cstr(text.c_str());
+        output = rb_utf8_str_new_cstr(text.c_str());
       } else {
         output = rb_ary_new();
         for (size_t i = 0; i < n_pieces; i++) {
@@ -249,7 +443,7 @@ private:
             rb_raise(rb_eSentencePieceError, "%s", status.message());
             return Qfalse;
           }
-          rb_ary_push(output, rb_str_new_cstr(text.c_str()));
+          rb_ary_push(output, rb_utf8_str_new_cstr(text.c_str()));
         }
       }
     } else {
@@ -265,7 +459,7 @@ private:
           rb_raise(rb_eSentencePieceError, "%s", status.message());
           return Qfalse;
         }
-        output = rb_str_new_cstr(text.c_str());
+        output = rb_utf8_str_new_cstr(text.c_str());
       } else {
         output = rb_ary_new();
         for (size_t i = 0; i < n_pieces; i++) {
@@ -282,7 +476,7 @@ private:
             rb_raise(rb_eSentencePieceError, "%s", status.message());
             return Qfalse;
           }
-          rb_ary_push(output, rb_str_new_cstr(text.c_str()));
+          rb_ary_push(output, rb_utf8_str_new_cstr(text.c_str()));
         }
       }
     }
@@ -290,6 +484,159 @@ private:
     return output;
   };
 
+  static VALUE _sentencepiece_processor_decode_pieces(VALUE self, VALUE pieces) {
+    if (!RB_TYPE_P(pieces, T_ARRAY)) {
+      rb_raise(rb_eArgError, "expected pieces to be an Array");
+      return Qnil;
+    }
+
+    std::vector<std::string> pcs;
+    const size_t n_pieces = RARRAY_LEN(pieces);
+    for (size_t i = 0; i < n_pieces; i++) {
+      VALUE et = rb_ary_entry(pieces, i);
+      pcs.push_back(StringValueCStr(et));
+    }
+
+    sentencepiece::SentencePieceProcessor* ptr = get_sentencepiece_processor(self);
+    const std::string text = ptr->DecodePieces(pcs);
+    VALUE output = rb_utf8_str_new_cstr(text.c_str());
+
+    return output;
+  };
+
+  static VALUE _sentencepiece_processor_decode_ids(VALUE self, VALUE ids) {
+    if (!RB_TYPE_P(ids, T_ARRAY)) {
+      rb_raise(rb_eArgError, "expected ids to be an Array");
+      return Qnil;
+    }
+
+    std::vector<int> pcs;
+    const size_t n_pieces = RARRAY_LEN(ids);
+    for (size_t i = 0; i < n_pieces; i++) {
+      VALUE et = rb_ary_entry(ids, i);
+      pcs.push_back(NUM2INT(et));
+    }
+
+    sentencepiece::SentencePieceProcessor* ptr = get_sentencepiece_processor(self);
+    const std::string text = ptr->DecodeIds(pcs);
+    VALUE output = rb_utf8_str_new_cstr(text.c_str());
+
+    return output;
+  };
+
+  static VALUE _sentencepiece_processor_encode_as_serialized_proto(VALUE self, VALUE text) {
+    if (!RB_TYPE_P(text, T_STRING)) {
+      rb_raise(rb_eArgError, "expected text to be a String");
+      return Qnil;
+    }
+
+    sentencepiece::SentencePieceProcessor* ptr = get_sentencepiece_processor(self);
+    const sentencepiece::util::bytes serialized = ptr->EncodeAsSerializedProto(StringValueCStr(text));
+    VALUE output = rb_str_new_cstr(serialized.c_str());
+
+    RB_GC_GUARD(text);
+    return output;
+  };
+
+  static VALUE _sentencepiece_processor_sample_encode_as_serialized_proto(int argc, VALUE* argv, VALUE self) {
+    VALUE kw_args = Qnil;
+    ID kw_table[2] = { rb_intern("nbest_size"), rb_intern("alpha") };
+    VALUE kw_values[2] = { Qundef, Qundef };
+
+    VALUE text = Qnil;
+    rb_scan_args(argc, argv, "1:", &text, &kw_args);
+    rb_get_kwargs(kw_args, kw_table, 2, 0, kw_values);
+
+    if (!RB_TYPE_P(text, T_STRING)) {
+      rb_raise(rb_eArgError, "expected text to be a String");
+      return Qnil;
+    }
+    if (!RB_INTEGER_TYPE_P(kw_values[0])) {
+      rb_raise(rb_eArgError, "expected nbest_size to be an Integer");
+      return Qnil;
+    }
+    if (!RB_INTEGER_TYPE_P(kw_values[1]) && !RB_FLOAT_TYPE_P(kw_values[1])) {
+      rb_raise(rb_eArgError, "expected alpha to be a Float");
+      return Qnil;
+    }
+
+    const int nbest_size = NUM2INT(kw_values[0]);
+    const float alpha = NUM2DBL(kw_values[1]);
+    sentencepiece::SentencePieceProcessor* ptr = get_sentencepiece_processor(self);
+    const sentencepiece::util::bytes serialized = ptr->SampleEncodeAsSerializedProto(StringValueCStr(text), nbest_size, alpha);
+    VALUE output = rb_str_new_cstr(serialized.c_str());
+
+    RB_GC_GUARD(text);
+    return output;
+  };
+
+  static VALUE _sentencepiece_processor_nbest_encode_as_serialized_proto(int argc, VALUE* argv, VALUE self) {
+    VALUE kw_args = Qnil;
+    ID kw_table[1] = { rb_intern("nbest_size") };
+    VALUE kw_values[1] = { Qundef };
+
+    VALUE text = Qnil;
+    rb_scan_args(argc, argv, "1:", &text, &kw_args);
+    rb_get_kwargs(kw_args, kw_table, 1, 0, kw_values);
+
+    if (!RB_TYPE_P(text, T_STRING)) {
+      rb_raise(rb_eArgError, "expected text to be a String");
+      return Qnil;
+    }
+    if (!RB_INTEGER_TYPE_P(kw_values[0])) {
+      rb_raise(rb_eArgError, "expected nbest_size to be an Integer");
+      return Qnil;
+    }
+
+    const int nbest_size = NUM2INT(kw_values[0]);
+    sentencepiece::SentencePieceProcessor* ptr = get_sentencepiece_processor(self);
+    const sentencepiece::util::bytes serialized = ptr->NBestEncodeAsSerializedProto(StringValueCStr(text), nbest_size);
+    VALUE output = rb_str_new_cstr(serialized.c_str());
+
+    RB_GC_GUARD(text);
+    return output;
+  };
+
+  static VALUE _sentencepiece_processor_decode_pieces_as_serialized_proto(VALUE self, VALUE pieces) {
+    if (!RB_TYPE_P(pieces, T_ARRAY)) {
+      rb_raise(rb_eArgError, "expected pieces to be an Array");
+      return Qnil;
+    }
+
+    std::vector<std::string> pcs;
+    const size_t n_pieces = RARRAY_LEN(pieces);
+    for (size_t i = 0; i < n_pieces; i++) {
+      VALUE et = rb_ary_entry(pieces, i);
+      pcs.push_back(StringValueCStr(et));
+    }
+
+    sentencepiece::SentencePieceProcessor* ptr = get_sentencepiece_processor(self);
+    const sentencepiece::util::bytes serialized = ptr->DecodePiecesAsSerializedProto(pcs);
+    VALUE output = rb_str_new_cstr(serialized.c_str());
+
+    return output;
+  };
+
+  static VALUE _sentencepiece_processor_decode_ids_as_serialized_proto(VALUE self, VALUE ids) {
+    if (!RB_TYPE_P(ids, T_ARRAY)) {
+      rb_raise(rb_eArgError, "expected ids to be an Array");
+      return Qnil;
+    }
+
+    std::vector<int> pcs;
+    const size_t n_pieces = RARRAY_LEN(ids);
+    for (size_t i = 0; i < n_pieces; i++) {
+      VALUE et = rb_ary_entry(ids, i);
+      pcs.push_back(NUM2INT(et));
+    }
+
+    sentencepiece::SentencePieceProcessor* ptr = get_sentencepiece_processor(self);
+    const sentencepiece::util::bytes serialized = ptr->DecodeIdsAsSerializedProto(pcs);
+    VALUE output = rb_str_new_cstr(serialized.c_str());
+
+    return output;
+  };
+
   static VALUE _sentencepiece_processor_piece_size(VALUE self) {
     sentencepiece::SentencePieceProcessor* ptr = get_sentencepiece_processor(self);
     return INT2NUM(ptr->GetPieceSize());
@@ -325,19 +672,49 @@ private:
     VALUE output = Qnil;
     sentencepiece::SentencePieceProcessor* ptr = get_sentencepiece_processor(self);
     if (RB_INTEGER_TYPE_P(ids)) {
-      const std::string piece = ptr->IdToPiece(NUM2INT(ids));
+      const int idx = NUM2INT(ids);
+      if (idx < 0 || idx >= ptr->GetPieceSize()) {
+        rb_raise(rb_eIndexError, "piece id %i is out of range", idx);
+        return Qnil;
+      }
+      const std::string piece = ptr->IdToPiece(idx);
       output = rb_utf8_str_new_cstr(piece.c_str());
     } else {
       const size_t n_ids = RARRAY_LEN(ids);
       output = rb_ary_new();
       for (size_t i = 0; i < n_ids; i++) {
         VALUE et = rb_ary_entry(ids, i);
-        const std::string piece = ptr->IdToPiece(NUM2INT(et));
+        const int idx = NUM2INT(et);
+        if (idx < 0 || idx >= ptr->GetPieceSize()) {
+          rb_raise(rb_eIndexError, "piece id %i is out of range", idx);
+          return Qnil;
+        }
+        const std::string piece = ptr->IdToPiece(idx);
         rb_ary_push(output, rb_utf8_str_new_cstr(piece.c_str()));
       }
     }
     return output;
   };
+
+  static VALUE _sentencepiece_processor_unk_id(VALUE self) {
+    sentencepiece::SentencePieceProcessor* ptr = get_sentencepiece_processor(self);
+    return INT2NUM(ptr->unk_id());
+  };
+
+  static VALUE _sentencepiece_processor_bos_id(VALUE self) {
+    sentencepiece::SentencePieceProcessor* ptr = get_sentencepiece_processor(self);
+    return INT2NUM(ptr->bos_id());
+  };
+
+  static VALUE _sentencepiece_processor_eos_id(VALUE self) {
+    sentencepiece::SentencePieceProcessor* ptr = get_sentencepiece_processor(self);
+    return INT2NUM(ptr->eos_id());
+  };
+
+  static VALUE _sentencepiece_processor_pad_id(VALUE self) {
+    sentencepiece::SentencePieceProcessor* ptr = get_sentencepiece_processor(self);
+    return INT2NUM(ptr->pad_id());
+  };
 };
 
 const rb_data_type_t RbSentencePieceProcessor::sentencepiece_processor_type = {
@@ -350,4 +727,42 @@ const rb_data_type_t RbSentencePieceProcessor::sentencepiece_processor_type = {
   RUBY_TYPED_FREE_IMMEDIATELY
 };
 
+namespace sentencepiece {
+class SentencePieceTrainerWrapper {
+public:
+  SentencePieceTrainerWrapper(){};
+  ~SentencePieceTrainerWrapper(){};
+
+  util::Status Train(absl::string_view args) {
+    return SentencePieceTrainer::Train(args);
+  };
+};
+} // namespace sentencepiece
+
+class RbSentencePieceTrainer {
+public:
+  static VALUE define_class(VALUE outer) {
+    rb_cSentencePieceTrainer = rb_define_class_under(outer, "SentencePieceTrainer", rb_cObject);
+    rb_define_singleton_method(rb_cSentencePieceTrainer, "train", RUBY_METHOD_FUNC(_sentencepiece_trainer_train), 1);
+    return rb_cSentencePieceTrainer;
+  };
+
+private:
+  static VALUE _sentencepiece_trainer_train(VALUE self, VALUE args) {
+    if (!RB_TYPE_P(args, T_STRING)) {
+      rb_raise(rb_eArgError, "expected args to be a String");
+      return Qnil;
+    }
+
+    const sentencepiece::util::Status status = sentencepiece::SentencePieceTrainer::Train(StringValueCStr(args));
+    if (!status.ok()) {
+      rb_raise(rb_eSentencePieceError, "%s", status.message());
+      return Qnil;
+    }
+
+    RB_GC_GUARD(args);
+    return Qnil;
+  };
+};
+
 #endif /* SENTENCEPIECE_HPP */

data/lib/sentencepiece/version.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+# sentencepiece.rb provides Ruby bindings for the SentencePiece.
 module SentencePiece
-  VERSION = '0.0.1'
+  # The version of sentencepiece.rb you install.
+  VERSION = '0.1.0'
 end

data/sig/sentencepiece.rbs

@@ -1,4 +1,63 @@
 module SentencePiece
   VERSION: String
-  # See the writing guide of rbs: https://github.com/ruby/rbs#guides
+
+  class SentencePieceTrainer
+    def self.train: (String args) -> void
+  end
+
+  class SentencePieceProcessor
+    def initialize: (?model_file: String model_file) -> void
+
+    def load: (String model_file) -> void
+
+    def encode: (String text, ?out_type: String out_type) -> (Array[Integer] | Array[String])
+              | (Array[String] text, ?out_type: String out_type) -> (Array[Array[Integer]] | Array[Array[String]])
+
+    def encode_as_ids: (String text) -> Array[Integer]
+
+    def encode_as_pieces: (String text) -> Array[String]
+
+    def encode_as_serialized_proto: (String text) -> String
+
+    def nbest_encode_as_ids: (String text, nbest_size: Integer nbest_size) -> Array[Array[Integer]]
+
+    def nbest_encode_as_pieces: (String text, nbest_size: Integer nbest_size) -> Array[Array[String]]
+
+    def nbest_encode_as_serialized_proto: (String text, nbest_size: Integer nbest_size) -> String
+
+    def sample_encode_as_ids: (String text, nbest_size: Integer nbest_size, alpha: Float alpha) -> Array[Integer]
+
+    def sample_encode_as_pieces: (String text, nbest_size: Integer nbest_size, alpha: Float alpha) -> Array[String]
+
+    def sample_encode_as_serialized_proto: (String text, nbest_size: Integer nbest_size, alpha: Float alpha) -> String
+
+    def decode: (Array[Integer], ?out_type: String out_type) -> String
+              | (Array[Array[Integer]], ?out_type: String out_type) -> Array[String]
+              | (Array[String], ?out_type: String out_type) -> String
+              | (Array[Array[String]], ?out_type: String out_type) -> Array[String]
+
+    def decode_ids: (Array[Integer]) -> String
+
+    def decode_ids_as_serialized_proto: (Array[Integer] ids) -> String
+
+    def decode_pieces: (Array[String]) -> String
+
+    def decode_pieces_as_serialized_proto: (Array[String] pieces) -> String
+
+    def id_to_piece: (Integer id) -> String
+                   | (Array[Integer] ids) -> Array[String]
+
+    def piece_to_id: (String piece) -> Integer
+                   | (Array[String] pieces) -> Array[Integer]
+
+    def piece_size: () -> Integer
+
+    def bos_id: () -> Integer
+
+    def eos_id: () -> Integer
+
+    def pad_id: () -> Integer
+
+    def unk_id: () -> Integer
+  end
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: sentencepiece
 version: !ruby/object:Gem::Version
-  version: 0.0.1
+  version: 0.1.0
 platform: ruby
 authors:
 - yoshoku
 autorequire:
 bindir: exe
 cert_chain: []
-date: 2023-03-21 00:00:00.000000000 Z
+date: 2023-03-26 00:00:00.000000000 Z
 dependencies: []
 description: |
   sentencepiece.rb provides Ruby bindings for the SentencePiece,
@@ -36,6 +36,7 @@ metadata:
   homepage_uri: https://github.com/yoshoku/sentencepiece.rb
   source_code_uri: https://github.com/yoshoku/sentencepiece.rb
   changelog_uri: https://github.com/yoshoku/sentencepiece.rb/blob/main/CHANGELOG.md
+  documentation_uri: https://yoshoku.github.io/sentencepiece.rb/doc/
   rubygems_mfa_required: 'true'
 post_install_message:
 rdoc_options: []