rlx_florence2/
generate.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! Autoregressive generation: greedy and beam search with the same logits
17//! processors HuggingFace applies for Florence-2 (`forced_bos_token_id`,
18//! `forced_eos_token_id`, `no_repeat_ngram_size`).
19
20use crate::config::Florence2TextConfig;
21use crate::runner::Florence2Model;
22use anyhow::Result;
23
24/// Generation knobs (defaults mirror Florence-2's `generation_config`).
25#[derive(Debug, Clone)]
26pub struct GenerateConfig {
27    pub max_new_tokens: usize,
28    pub num_beams: usize,
29    pub length_penalty: f64,
30    pub early_stopping: bool,
31}
32
33impl GenerateConfig {
34    pub fn from_text(t: &Florence2TextConfig, max_new_tokens: usize) -> Self {
35        Self {
36            max_new_tokens,
37            num_beams: t.num_beams,
38            length_penalty: 1.0,
39            early_stopping: true,
40        }
41    }
42}
43
44fn log_softmax(logits: &[f32]) -> Vec<f32> {
45    let max = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
46    let mut sum = 0f64;
47    for &l in logits {
48        sum += ((l - max) as f64).exp();
49    }
50    let lse = max as f64 + sum.ln();
51    logits.iter().map(|&l| (l as f64 - lse) as f32).collect()
52}
53
54/// Apply Florence-2's logits processors in place (HF order: no-repeat-ngram,
55/// forced-bos, forced-eos). `cur_len` is the current decoder length, i.e. the
56/// position index of the token about to be produced.
57fn apply_processors(
58    logits: &mut [f32],
59    seq: &[u32],
60    cur_len: usize,
61    max_len: usize,
62    t: &Florence2TextConfig,
63    no_repeat_ngram: usize,
64) {
65    // no_repeat_ngram_size: ban tokens completing a previously seen n-gram.
66    if no_repeat_ngram > 0 && seq.len() >= no_repeat_ngram {
67        let n = no_repeat_ngram;
68        let prefix = &seq[seq.len() + 1 - n..]; // last (n-1) tokens
69        for i in 0..=seq.len() - n {
70            if seq[i..i + n - 1] == *prefix {
71                let banned = seq[i + n - 1] as usize;
72                if banned < logits.len() {
73                    logits[banned] = f32::NEG_INFINITY;
74                }
75            }
76        }
77    }
78    // forced_bos_token_id: first generated token (position 1).
79    if cur_len == 1 {
80        force_token(logits, t.forced_bos_token_id);
81    }
82    // forced_eos_token_id: last position.
83    if cur_len == max_len - 1 {
84        force_token(logits, t.forced_eos_token_id);
85    }
86}
87
88fn force_token(logits: &mut [f32], token: u32) {
89    for (i, l) in logits.iter_mut().enumerate() {
90        *l = if i as u32 == token {
91            0.0
92        } else {
93            f32::NEG_INFINITY
94        };
95    }
96}
97
98fn argmax(a: &[f32]) -> u32 {
99    let mut bi = 0u32;
100    let mut bv = f32::NEG_INFINITY;
101    for (i, &v) in a.iter().enumerate() {
102        if v > bv {
103            bv = v;
104            bi = i as u32;
105        }
106    }
107    bi
108}
109
110impl Florence2Model {
111    /// Greedy decode. `encoder_hidden` is `[enc_seq · d]`. Returns the full
112    /// token sequence including the start token (and trailing EOS).
113    pub fn generate_greedy(
114        &mut self,
115        encoder_hidden: &[f32],
116        enc_seq: usize,
117        opts: &GenerateConfig,
118    ) -> Result<Vec<u32>> {
119        let t = self.config().text.clone();
120        let max_len = opts.max_new_tokens + 1;
121        let no_ngram = t.no_repeat_ngram_size;
122        let mut seq = vec![t.decoder_start_token_id];
123        while seq.len() < max_len {
124            let mut last = self.decode_logits(&seq, encoder_hidden, enc_seq, max_len)?;
125            apply_processors(&mut last, &seq, seq.len(), max_len, &t, no_ngram);
126            let next = argmax(&last);
127            seq.push(next);
128            if next == t.eos_token_id {
129                break;
130            }
131        }
132        Ok(seq)
133    }
134
135    /// Beam search (HF semantics: log-prob accumulation, `2·num_beams`
136    /// candidate pool, length-penalty normalization, early stopping).
137    pub fn generate_beam(
138        &mut self,
139        encoder_hidden: &[f32],
140        enc_seq: usize,
141        opts: &GenerateConfig,
142    ) -> Result<Vec<u32>> {
143        let t = self.config().text.clone();
144        let nb = opts.num_beams.max(1);
145        if nb == 1 {
146            return self.generate_greedy(encoder_hidden, enc_seq, opts);
147        }
148        let max_len = opts.max_new_tokens + 1;
149        let no_ngram = t.no_repeat_ngram_size;
150        let eos = t.eos_token_id;
151
152        // Active beams: (tokens, cumulative_logprob). Seed only beam 0 so the
153        // first expansion does not produce `num_beams` identical hypotheses.
154        let mut beams: Vec<(Vec<u32>, f64)> = vec![(vec![t.decoder_start_token_id], 0.0)];
155        let mut finished: Vec<(Vec<u32>, f64)> = Vec::new();
156
157        for _ in 0..max_len - 1 {
158            let mut cands: Vec<(Vec<u32>, f64)> = Vec::new();
159            for (seq, score) in &beams {
160                let mut last = self.decode_logits(seq, encoder_hidden, enc_seq, max_len)?;
161                apply_processors(&mut last, seq, seq.len(), max_len, &t, no_ngram);
162                let logp = log_softmax(&last);
163                // Top 2*nb tokens for this beam.
164                let mut idx: Vec<usize> = (0..logp.len()).collect();
165                idx.select_nth_unstable_by(2 * nb, |&a, &b| logp[b].partial_cmp(&logp[a]).unwrap());
166                for &tok in idx.iter().take(2 * nb) {
167                    if logp[tok] == f32::NEG_INFINITY {
168                        continue;
169                    }
170                    let mut ns = seq.clone();
171                    ns.push(tok as u32);
172                    cands.push((ns, score + logp[tok] as f64));
173                }
174            }
175            // Keep the best `2*nb` candidates, route EOS to finished.
176            cands.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
177            let mut next_beams: Vec<(Vec<u32>, f64)> = Vec::new();
178            for (seq, score) in cands {
179                if *seq.last().unwrap() == eos {
180                    let norm = score / (seq.len() as f64).powf(opts.length_penalty);
181                    finished.push((seq, norm));
182                } else {
183                    next_beams.push((seq, score));
184                }
185                if next_beams.len() >= nb {
186                    break;
187                }
188            }
189            beams = next_beams;
190            if beams.is_empty() {
191                break;
192            }
193            // Early stopping: enough finished hypotheses and the best active
194            // beam cannot beat them.
195            if opts.early_stopping && finished.len() >= nb {
196                break;
197            }
198        }
199
200        // Add any unfinished beams as candidates (normalized).
201        for (seq, score) in beams {
202            let norm = score / (seq.len() as f64).powf(opts.length_penalty);
203            finished.push((seq, norm));
204        }
205        finished.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
206        Ok(finished
207            .into_iter()
208            .next()
209            .map(|(s, _)| s)
210            .unwrap_or_default())
211    }
212}
rlx_florence2/generate.rs

rlx_florence2/
generate.rs
