add log probs

Author: minhthuc2502

include/ctranslate2/decoding.h

@@ -15,14 +15,16 @@ namespace ctranslate2 {
     std::vector<std::vector<size_t>> hypotheses;
     std::vector<float> scores;
     std::vector<std::vector<std::vector<float>>> attention;
+    std::vector<std::vector<StorageView>> log_probs_vocab;
   };
 
   struct DecodingStepResult {
     size_t step;
     size_t batch_id;
     size_t token_id;
     size_t hypothesis_id;
-    std::optional<float> log_prob;
+    std::optional<float> score;
+    std::optional<StorageView> log_probs;
     bool is_last = false;
   };
 
@@ -41,6 +43,7 @@ namespace ctranslate2 {
            const dim_t min_length,
            const bool return_scores = false,
            const bool return_attention = false,
+           const bool return_log_probs_vocab = true,
            const bool return_prefix = true,
            const size_t num_hypotheses = 1,
            const bool include_eos_in_hypotheses = true,
@@ -67,6 +70,7 @@ namespace ctranslate2 {
            const dim_t min_length,
            const bool return_scores = false,
            const bool return_attention = false,
+           const bool return_log_probs_vocab = true,
            const bool return_prefix = true,
            const size_t num_hypotheses = 1,
            const bool include_eos_in_hypotheses = true,
@@ -118,6 +122,7 @@ namespace ctranslate2 {
            const dim_t min_length,
            const bool return_scores = false,
            const bool return_attention = false,
+           const bool return_log_probs_vocab = true,
            const bool return_prefix = true,
            const size_t num_hypotheses = 1,
            const bool include_eos_in_hypotheses = true,
@@ -149,6 +154,7 @@ namespace ctranslate2 {
     bool include_eos_in_hypotheses = true;
     bool return_scores = false;
     bool return_attention = false;
+    bool return_log_probs_vocab = false;
     bool return_alternatives = false;
     bool return_prefix = true;
     float min_alternative_expansion_prob = 0;

include/ctranslate2/generation.h

@@ -53,6 +53,8 @@ namespace ctranslate2 {
 
     // Include scores in the result.
     bool return_scores = false;
+    // Include log probs of each token in the result
+    bool return_log_probs_vocab = false;
 
     // Return alternatives at the first unconstrained decoding position. This is typically
     // used with a prefix to provide alternatives at a specifc location.
@@ -79,6 +81,7 @@ namespace ctranslate2 {
     std::vector<std::vector<std::string>> sequences;
     std::vector<std::vector<size_t>> sequences_ids;
     std::vector<float> scores;
+    std::vector<std::vector<StorageView>> log_probs;
 
     size_t num_sequences() const {
       return sequences.size();
@@ -95,7 +98,8 @@ namespace ctranslate2 {
     size_t token_id;
     size_t hypothesis_id;
     std::string token;
-    std::optional<float> log_prob;
+    std::optional<float> score;
+    std::optional<StorageView> log_probs;
     bool is_last;
 
     GenerationStepResult() = default;
@@ -105,7 +109,8 @@ namespace ctranslate2 {
       , token_id(result.token_id)
       , hypothesis_id(result.hypothesis_id)
       , token(vocabulary.to_token(result.token_id))
-      , log_prob(result.log_prob)
+      , score(result.score)
+      , log_probs(result.log_probs)
       , is_last(result.is_last)
     {
     }

include/ctranslate2/models/whisper.h

@@ -41,6 +41,9 @@ namespace ctranslate2 {
       // Include scores in the result.
       bool return_scores = false;
 
+      // Include log probs of each token in the result
+      bool return_log_probs_vocab = false;
+
       // Include the probability of the no speech token in the result.
       bool return_no_speech_prob = false;
 
@@ -59,6 +62,7 @@ namespace ctranslate2 {
       std::vector<std::vector<std::string>> sequences;
       std::vector<std::vector<size_t>> sequences_ids;
       std::vector<float> scores;
+      std::vector<std::vector<StorageView>> log_probs;
       float no_speech_prob = 0;
 
       size_t num_sequences() const {

include/ctranslate2/translation.h

@@ -67,6 +67,8 @@ namespace ctranslate2 {
     bool return_scores = false;
     // Store attention vectors in the TranslationResult class.
     bool return_attention = false;
+    // Store log probs matrix in the TranslationResult class.
+    bool return_log_probs_vocab = false;
 
     // Return alternatives at the first unconstrained decoding position. This is typically
     // used with a target prefix to provide alternatives at a specifc location in the
@@ -87,6 +89,7 @@ namespace ctranslate2 {
     std::vector<std::vector<std::string>> hypotheses;
     std::vector<float> scores;
     std::vector<std::vector<std::vector<float>>> attention;
+    std::vector<std::vector<StorageView>> log_probs;
 
     TranslationResult(std::vector<std::vector<std::string>> hypotheses_)
       : hypotheses(std::move(hypotheses_))
@@ -95,10 +98,12 @@ namespace ctranslate2 {
 
     TranslationResult(std::vector<std::vector<std::string>> hypotheses_,
                       std::vector<float> scores_,
-                      std::vector<std::vector<std::vector<float>>> attention_)
+                      std::vector<std::vector<std::vector<float>>> attention_,
+                      std::vector<std::vector<StorageView>> log_probs_)
       : hypotheses(std::move(hypotheses_))
       , scores(std::move(scores_))
       , attention(std::move(attention_))
+      , log_probs(std::move(log_probs_))
     {
     }
 

python/cpp/generator.cc

@@ -33,6 +33,7 @@ namespace ctranslate2 {
                      bool cache_static_prompt,
                      bool include_prompt_in_result,
                      bool return_scores,
+                     bool return_log_probs_vocab,
                      bool return_alternatives,
                      float min_alternative_expansion_prob,
                      size_t sampling_topk,
@@ -58,6 +59,7 @@ namespace ctranslate2 {
         options.num_hypotheses = num_hypotheses;
         options.return_end_token = return_end_token;
         options.return_scores = return_scores;
+        options.return_log_probs_vocab = return_log_probs_vocab;
         options.return_alternatives = return_alternatives;
         options.cache_static_prompt = cache_static_prompt;
         options.include_prompt_in_result = include_prompt_in_result;
@@ -203,6 +205,7 @@ namespace ctranslate2 {
              py::arg("cache_static_prompt")=true,
              py::arg("include_prompt_in_result")=true,
              py::arg("return_scores")=false,
+             py::arg("return_log_probs_vocab")=false,
              py::arg("return_alternatives")=false,
              py::arg("min_alternative_expansion_prob")=0,
              py::arg("sampling_topk")=1,
@@ -260,6 +263,7 @@ namespace ctranslate2 {
                      reuse it for future generations using the same static prompt.
                    include_prompt_in_result: Include the :obj:`start_tokens` in the result.
                    return_scores: Include the scores in the output.
+                   return_log_probs_vocab: Include log probs for each token in the output
                    return_alternatives: Return alternatives at the first unconstrained decoding position.
                    min_alternative_expansion_prob: Minimum initial probability to expand an alternative.
                    sampling_topk: Randomly sample predictions from the top K candidates.

python/cpp/translator.cc

@@ -52,6 +52,7 @@ namespace ctranslate2 {
                      size_t min_decoding_length,
                      bool use_vmap,
                      bool with_scores,
+                     bool return_log_probs_vocab,
                      size_t sampling_topk,
                      float sampling_topp,
                      float sampling_temperature,
@@ -141,6 +142,7 @@ namespace ctranslate2 {
                       size_t min_decoding_length,
                       bool use_vmap,
                       bool return_scores,
+                      bool return_log_probs_vocab,
                       bool return_attention,
                       bool return_alternatives,
                       float min_alternative_expansion_prob,
@@ -172,6 +174,7 @@ namespace ctranslate2 {
         options.use_vmap = use_vmap;
         options.return_end_token = return_end_token;
         options.return_scores = return_scores;
+        options.return_log_probs_vocab = return_log_probs_vocab;
         options.return_attention = return_attention;
         options.return_alternatives = return_alternatives;
         options.min_alternative_expansion_prob = min_alternative_expansion_prob;
@@ -354,6 +357,7 @@ namespace ctranslate2 {
              py::arg("min_decoding_length")=1,
              py::arg("use_vmap")=false,
              py::arg("return_scores")=false,
+             py::arg("return_log_probs_vocab")=false,
              py::arg("return_attention")=false,
              py::arg("return_alternatives")=false,
              py::arg("min_alternative_expansion_prob")=0,
@@ -396,6 +400,7 @@ namespace ctranslate2 {
                    min_decoding_length: Minimum prediction length.
                    use_vmap: Use the vocabulary mapping file saved in this model
                    return_scores: Include the scores in the output.
+                   return_log_probs_vocab: Include the log probs of each token in the output
                    return_attention: Include the attention vectors in the output.
                    return_alternatives: Return alternatives at the first unconstrained decoding position.
                    min_alternative_expansion_prob: Minimum initial probability to expand an alternative.

python/cpp/whisper.cc

@@ -40,6 +40,7 @@ namespace ctranslate2 {
                size_t no_repeat_ngram_size,
                size_t max_length,
                bool return_scores,
+               bool return_log_probs_vocab,
                bool return_no_speech_prob,
                size_t max_initial_timestamp_index,
                bool suppress_blank,
@@ -59,6 +60,7 @@ namespace ctranslate2 {
         options.max_length = max_length;
         options.num_hypotheses = num_hypotheses;
         options.return_scores = return_scores;
+        options.return_log_probs_vocab = return_log_probs_vocab;
         options.return_no_speech_prob = return_no_speech_prob;
         options.max_initial_timestamp_index = max_initial_timestamp_index;
         options.suppress_blank = suppress_blank;
@@ -247,6 +249,7 @@ namespace ctranslate2 {
              py::arg("no_repeat_ngram_size")=0,
              py::arg("max_length")=448,
              py::arg("return_scores")=false,
+             py::arg("return_log_probs_vocab")=false,
              py::arg("return_no_speech_prob")=false,
              py::arg("max_initial_timestamp_index")=50,
              py::arg("suppress_blank")=true,
@@ -276,6 +279,7 @@ namespace ctranslate2 {
                      (set 0 to disable).
                    max_length: Maximum generation length.
                    return_scores: Include the scores in the output.
+                   return_log_probs_vocab: Include the log probs in the output
                    return_no_speech_prob: Include the probability of the no speech token in the
                      result.
                    max_initial_timestamp_index: Maximum index of the first predicted timestamp.