ct2rs 0.9.18

Rust bindings for OpenNMT/CTranslate2
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

// whisper.cpp
//
// Copyright (c) 2023-2024 Junpei Kawamoto
//
// This software is released under the MIT License.
//
// http://opensource.org/licenses/mit-license.php

#include <utility>

#include "ct2rs/src/sys/whisper.rs.h"

#include "ct2rs/include/types.h"

using rust::Slice;
using rust::Vec;

std::unique_ptr<StorageView> Whisper::encode(
    const StorageView& features,
    const bool to_cpu
) const {
    auto future = impl->encode(
        features,
        to_cpu
    );

    StorageView storage_view = future.get();
    return std::make_unique<StorageView>(std::move(storage_view));
}

Vec<WhisperGenerationResult> Whisper::generate(
    const StorageView& features,
    const Slice<const VecStr> prompts,
    const WhisperOptions& opts
) const {
    auto futures = impl->generate(
        features,
        from_rust(prompts),
        ctranslate2::models::WhisperOptions {
            opts.beam_size,
            opts.patience,
            opts.length_penalty,
            opts.repetition_penalty,
            opts.no_repeat_ngram_size,
            opts.max_length,
            opts.sampling_topk,
            opts.sampling_temperature,
            opts.num_hypotheses,
            opts.return_scores,
            opts.return_logits_vocab,
            opts.return_no_speech_prob,
            opts.max_initial_timestamp_index,
            opts.suppress_blank,
            from_rust(opts.suppress_tokens),
        }
    );

    Vec<WhisperGenerationResult> res;
    for (auto& future : futures) {
        const auto& r = future.get();
        res.push_back(WhisperGenerationResult {
            to_rust<VecString>(r.sequences),
            to_rust<VecUSize>(r.sequences_ids),
            to_rust(r.scores),
            r.no_speech_prob,
        });
    }

    return res;
}

Vec<VecDetectionResult> Whisper::detect_language(const StorageView& features) const {
    auto futures = impl->detect_language(features);

    Vec<VecDetectionResult> res;
    for (auto& future : futures) {
        const auto& r = future.get();

        Vec<DetectionResult> pairs;
        for (auto& pair : r) {
            pairs.push_back(DetectionResult {
                std::get<0>(pair),
                std::get<1>(pair),
            });
        }

        res.push_back(VecDetectionResult { pairs });
    }

    return res;
}

Vec<WhisperAlignmentResult> Whisper::align(
    const StorageView& features,
    const Slice<const size_t> start_sequence,
    const Slice<const Vec<size_t>> text_tokens,
    const Slice<const size_t> num_frames,
    int64_t median_filter_width
) const {
    std::vector<size_t> start_sequence_cxx(start_sequence.begin(), start_sequence.end());

    std::vector<std::vector<size_t>> text_tokens_cxx;
    for (auto& seq_text_tokens : text_tokens) {
        text_tokens_cxx.emplace_back(seq_text_tokens.begin(), seq_text_tokens.end());
    }

    std::vector<size_t> num_frames_cxx(num_frames.begin(), num_frames.end());

    auto futures = impl->align(
        features, start_sequence_cxx, text_tokens_cxx, num_frames_cxx, median_filter_width
    );

    Vec<WhisperAlignmentResult> res;
    for (auto& future : futures) {
        const auto& result_cxx = future.get();

        WhisperAlignmentResult result;
        for (auto& token_alignment_cxx : result_cxx.alignments) {
            result.alignments.push_back(
                WhisperTokenAlignment { token_alignment_cxx.first, token_alignment_cxx.second }
            );
        }
        result.text_token_probs = to_rust(result_cxx.text_token_probs);

        res.push_back(result);
    }

    return res;
}