lib.rs 51.9 KB
Newer Older
1
/// Text Generation Inference Webserver
OlivierDehaene's avatar
OlivierDehaene committed
2
pub mod config;
Nicolas Patry's avatar
Nicolas Patry committed
3
pub mod infer;
Olivier Dehaene's avatar
Olivier Dehaene committed
4
pub mod server;
Nicolas Patry's avatar
Nicolas Patry committed
5
pub mod validation;
Olivier Dehaene's avatar
Olivier Dehaene committed
6

7
8
#[cfg(feature = "kserve")]
mod kserve;
Nicolas Patry's avatar
Nicolas Patry committed
9
pub mod logging;
10

11
mod sagemaker;
12
pub mod usage_stats;
Nicolas Patry's avatar
Nicolas Patry committed
13
mod vertex;
14

Nicolas Patry's avatar
Nicolas Patry committed
15
16
use crate::infer::{Infer, InferError};
use crate::server::prepare_chat_input;
17
18
use pyo3::prelude::*;
use pyo3::types::IntoPyDict;
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
19
use serde::{Deserialize, Serialize};
20
use tokenizers::Encoding;
Nicolas Patry's avatar
Nicolas Patry committed
21
use tracing::warn;
22
use utoipa::ToSchema;
Olivier Dehaene's avatar
Olivier Dehaene committed
23
use validation::Validation;
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
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
#[derive(Clone)]
pub enum Tokenizer {
    Python {
        tokenizer_name: String,
        revision: Option<String>,
    },
    Rust(tokenizers::Tokenizer),
}

pub struct PyTokenizer<'a>(pyo3::Bound<'a, pyo3::PyAny>);

impl<'a> PyTokenizer<'a> {
    fn from_py(
        py: Python<'a>,
        tokenizer_name: String,
        revision: Option<String>,
    ) -> PyResult<PyTokenizer<'a>> {
        let transformers = py.import_bound("transformers")?;
        let auto = transformers.getattr("AutoTokenizer")?;
        let from_pretrained = auto.getattr("from_pretrained")?;
        let args = (tokenizer_name,);
        let kwargs = if let Some(rev) = &revision {
            [("revision", rev.to_string())].into_py_dict_bound(py)
        } else {
            pyo3::types::PyDict::new_bound(py)
        };
        let tokenizer = from_pretrained.call(args, Some(&kwargs))?;
        tracing::info!("Loaded a python tokenizer");
        Ok(PyTokenizer(tokenizer))
    }
}

trait TokenizerTrait {
    fn encode_trait(
        &self,
        query: String,
        add_special_tokens: bool,
    ) -> Result<tokenizers::Encoding, Box<dyn std::error::Error + Send + Sync>>;
}

impl TokenizerTrait for tokenizers::Tokenizer {
    fn encode_trait(
        &self,
        query: String,
        add_special_tokens: bool,
    ) -> Result<tokenizers::Encoding, Box<dyn std::error::Error + Send + Sync>> {
        self.encode(query, add_special_tokens)
    }
}

impl<'a> TokenizerTrait for PyTokenizer<'a> {
    fn encode_trait(
        &self,
        query: String,
        add_special_tokens: bool,
    ) -> Result<tokenizers::Encoding, Box<dyn std::error::Error + Send + Sync>> {
        let py = self.0.py();
        let kwargs = [
            ("text", query.into_py(py)),
            ("add_special_tokens", add_special_tokens.into_py(py)),
        ]
        .into_py_dict_bound(py);
        let encode = self.0.getattr("encode")?;
        let input_ids: Vec<u32> = encode.call((), Some(&kwargs))?.extract()?;
        Ok(Encoding::new(
            input_ids,
            vec![],                           // type ids
            vec![],                           // tokens (strings)
            vec![],                           // words
            vec![],                           // offsets
            vec![],                           // special_tokens_mask
            vec![],                           // attention_mask
            vec![],                           // overflowing
            std::collections::HashMap::new(), //sequence_ranges
        ))
    }
}

103
104
/// Hub type
#[derive(Clone, Debug, Deserialize)]
105
pub struct HubModelInfo {
106
107
108
109
110
111
    #[serde(rename(deserialize = "id"))]
    pub model_id: String,
    pub sha: Option<String>,
    pub pipeline_tag: Option<String>,
}

112
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
113
114
115
116
117
pub struct ChatTemplate {
    name: String,
    template: String,
}

118
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
119
120
121
122
123
124
#[serde(untagged)]
pub enum ChatTemplateVersions {
    Single(String),
    Multiple(Vec<ChatTemplate>),
}

125
126
use std::path::Path;

127
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
128
pub struct HubTokenizerConfig {
129
    pub chat_template: Option<ChatTemplateVersions>,
130
    pub completion_template: Option<String>,
131
132
    pub bos_token: Option<TokenizerConfigToken>,
    pub eos_token: Option<TokenizerConfigToken>,
133
134
135
    pub tokenizer_class: Option<String>,
    pub add_bos_token: Option<bool>,
    pub add_eos_token: Option<bool>,
136
137
138
}

impl HubTokenizerConfig {
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
    pub fn from_file<P: AsRef<Path>>(filename: P) -> Option<Self> {
        std::fs::read_to_string(filename)
            .ok()
            .and_then(|content| serde_json::from_str(&content).ok())
    }
}

#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
#[serde(untagged)]
pub enum TokenizerConfigToken {
    String(String),
    Object { content: String },
}

impl TokenizerConfigToken {
    pub fn as_str(&self) -> &str {
        match self {
            TokenizerConfigToken::String(s) => s,
            TokenizerConfigToken::Object { content } => content,
        }
159
160
161
    }
}

162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "processor_class")]
pub enum HubPreprocessorConfig {
    Idefics2Processor(Idefics2Preprocessor),
}

impl HubPreprocessorConfig {
    pub fn from_file<P: AsRef<std::path::Path>>(filename: P) -> Option<Self> {
        let content = std::fs::read_to_string(filename).ok()?;
        serde_json::from_str(&content).ok()
    }
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Idefics2Preprocessor {
    #[serde(default)]
    do_image_splitting: bool,
}

drbh's avatar
drbh committed
181
182
183
184
185
186
187
188
#[derive(Debug, Clone, Deserialize, Default)]
pub struct HubProcessorConfig {
    pub chat_template: Option<ChatTemplateVersions>,
    pub image_seq_len: usize,
    pub processor_class: Option<String>,
}

impl HubProcessorConfig {
189
190
191
192
    pub fn from_file<P: AsRef<Path>>(filename: P) -> Option<Self> {
        std::fs::read_to_string(filename)
            .ok()
            .and_then(|content| serde_json::from_str(&content).ok())
drbh's avatar
drbh committed
193
194
195
    }
}

196
#[derive(Clone, Debug, Deserialize, ToSchema, Serialize)]
Nicolas Patry's avatar
Nicolas Patry committed
197
#[cfg_attr(test, derive(PartialEq))]
drbh's avatar
drbh committed
198
199
#[serde(tag = "type", content = "value")]
pub(crate) enum GrammarType {
200
201
202
203
204
    /// A string that represents a [JSON Schema](https://json-schema.org/).
    ///
    /// JSON Schema is a declarative language that allows to annotate JSON documents
    /// with types and descriptions.
    #[serde(rename = "json")]
drbh's avatar
drbh committed
205
    #[serde(alias = "json_object")]
206
207
    #[schema(example = json ! ({"properties": {"location":{"type": "string"}}}))]
    Json(serde_json::Value),
drbh's avatar
drbh committed
208
209
210
211
    #[serde(rename = "regex")]
    Regex(String),
}

212
213
#[derive(Clone, Debug, Serialize, ToSchema)]
pub struct Info {
214
    /// Model info
215
216
217
218
    #[schema(example = "bigscience/blomm-560m")]
    pub model_id: String,
    #[schema(nullable = true, example = "e985a63cdc139290c5f700ff1929f0b5942cced2")]
    pub model_sha: Option<String>,
Nicolas Patry's avatar
Nicolas Patry committed
219
220
221
222
    // #[schema(example = "torch.float16")]
    // pub model_dtype: String,
    // #[schema(example = "cuda")]
    // pub model_device_type: String,
223
224
    #[schema(nullable = true, example = "text-generation")]
    pub model_pipeline_tag: Option<String>,
Nicolas Patry's avatar
Nicolas Patry committed
225

226
227
228
229
230
231
232
233
    /// Router Parameters
    #[schema(example = "128")]
    pub max_concurrent_requests: usize,
    #[schema(example = "2")]
    pub max_best_of: usize,
    #[schema(example = "4")]
    pub max_stop_sequences: usize,
    #[schema(example = "1024")]
OlivierDehaene's avatar
OlivierDehaene committed
234
    pub max_input_tokens: usize,
235
236
237
238
    #[schema(example = "2048")]
    pub max_total_tokens: usize,
    #[schema(example = "2")]
    pub validation_workers: usize,
239
240
    #[schema(example = "32")]
    pub max_client_batch_size: usize,
Nicolas Patry's avatar
Nicolas Patry committed
241

242
    /// Router Info
243
244
    #[schema(example = "text-generation-router")]
    pub router: &'static str,
245
246
247
248
    #[schema(example = "0.5.0")]
    pub version: &'static str,
    #[schema(nullable = true, example = "null")]
    pub sha: Option<&'static str>,
249
250
    #[schema(nullable = true, example = "null")]
    pub docker_label: Option<&'static str>,
251
252
}

drbh's avatar
drbh committed
253
#[derive(Clone, Debug, Deserialize, ToSchema, Default)]
Nicolas Patry's avatar
Nicolas Patry committed
254
#[cfg_attr(test, derive(PartialEq))]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
255
pub(crate) struct GenerateParameters {
256
    /// Generate best_of sequences and return the one if the highest token logprobs.
257
258
259
    #[serde(default)]
    #[schema(exclusive_minimum = 0, nullable = true, default = "null", example = 1)]
    pub best_of: Option<usize>,
260
261

    /// The value used to module the logits distribution.
262
263
264
265
266
267
268
269
    #[serde(default)]
    #[schema(
        exclusive_minimum = 0.0,
        nullable = true,
        default = "null",
        example = 0.5
    )]
    pub temperature: Option<f32>,
270
271
272

    /// The parameter for repetition penalty. 1.0 means no penalty.
    /// See [this paper](https://arxiv.org/pdf/1909.05858.pdf) for more details.
273
274
275
276
277
278
279
280
    #[serde(default)]
    #[schema(
        exclusive_minimum = 0.0,
        nullable = true,
        default = "null",
        example = 1.03
    )]
    pub repetition_penalty: Option<f32>,
281
282
283
284

    /// The parameter for frequency penalty. 1.0 means no penalty
    /// Penalize new tokens based on their existing frequency in the text so far,
    /// decreasing the model's likelihood to repeat the same line verbatim.
285
    #[serde(default)]
286
287
288
289
290
291
292
    #[schema(
        exclusive_minimum = -2.0,
        nullable = true,
        default = "null",
        example = 0.1
    )]
    pub frequency_penalty: Option<f32>,
293
294

    /// The number of highest probability vocabulary tokens to keep for top-k-filtering.
295
    #[serde(default)]
296
297
    #[schema(exclusive_minimum = 0, nullable = true, default = "null", example = 10)]
    pub top_k: Option<i32>,
298
299

    /// Top-p value for nucleus sampling.
300
301
302
303
304
305
306
307
308
    #[serde(default)]
    #[schema(
        exclusive_minimum = 0.0,
        maximum = 1.0,
        nullable = true,
        default = "null",
        example = 0.95
    )]
    pub top_p: Option<f32>,
309
310
311

    /// Typical Decoding mass
    /// See [Typical Decoding for Natural Language Generation](https://arxiv.org/abs/2202.00666) for more information.
312
    #[serde(default)]
313
314
315
316
317
318
319
320
    #[schema(
        exclusive_minimum = 0.0,
        maximum = 1.0,
        nullable = true,
        default = "null",
        example = 0.95
    )]
    pub typical_p: Option<f32>,
321
322

    /// Activate logits sampling.
323
    #[serde(default)]
324
    #[schema(default = "false", example = true)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
325
    pub do_sample: bool,
326
327

    /// Maximum number of tokens to generate.
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
328
    #[serde(default = "default_max_new_tokens")]
329
    #[schema(nullable = true, default = "100", example = "20")]
330
    pub max_new_tokens: Option<u32>,
331
332

    /// Whether to prepend the prompt to the generated text
OlivierDehaene's avatar
OlivierDehaene committed
333
    #[serde(default)]
334
    #[schema(nullable = true, default = "null", example = false)]
335
    pub return_full_text: Option<bool>,
336
337

    /// Stop generating tokens if a member of `stop` is generated.
338
    #[serde(default)]
339
    #[schema(inline, max_items = 4, example = json ! (["photographer"]))]
340
    pub stop: Vec<String>,
341
342

    /// Truncate inputs tokens to the given size.
OlivierDehaene's avatar
OlivierDehaene committed
343
    #[serde(default)]
344
    #[schema(nullable = true, default = "null", example = "null")]
345
    pub truncate: Option<usize>,
346
347

    /// Watermarking with [A Watermark for Large Language Models](https://arxiv.org/abs/2301.10226).
348
    #[serde(default)]
349
350
    #[schema(default = "false", example = true)]
    pub watermark: bool,
351
352

    /// Whether to return generation details.
353
    #[serde(default)]
354
    #[schema(default = "true")]
OlivierDehaene's avatar
OlivierDehaene committed
355
    pub details: bool,
356
357

    /// Whether to return decoder input token logprobs and ids.
358
    #[serde(default)]
359
    #[schema(default = "false")]
360
    pub decoder_input_details: bool,
361
362

    /// Random sampling seed.
363
    #[serde(default)]
364
365
366
367
368
369
    #[schema(
        exclusive_minimum = 0,
        nullable = true,
        default = "null",
        example = "null"
    )]
370
    pub seed: Option<u64>,
371
372

    /// The number of highest probability vocabulary tokens to keep for top-n-filtering.
Nicolas Patry's avatar
Nicolas Patry committed
373
374
375
    #[serde(default)]
    #[schema(exclusive_minimum = 0, nullable = true, default = "null", example = 5)]
    pub top_n_tokens: Option<u32>,
376
377

    /// Grammar constraints for the generation.
drbh's avatar
drbh committed
378
    #[serde(default)]
379
    #[schema(nullable = true, default = "null", example = "null")]
drbh's avatar
drbh committed
380
    pub grammar: Option<GrammarType>,
drbh's avatar
drbh committed
381
382
383
384
385

    /// Lora adapter id
    #[serde(default)]
    #[schema(nullable = true, default = "null", example = "null")]
    pub adapter_id: Option<String>,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
386
387
}

388
fn default_max_new_tokens() -> Option<u32> {
389
    Some(100)
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
390
391
392
393
}

fn default_parameters() -> GenerateParameters {
    GenerateParameters {
394
        best_of: None,
395
396
        temperature: None,
        repetition_penalty: None,
397
        frequency_penalty: None,
398
399
        top_k: None,
        top_p: None,
400
        typical_p: None,
401
        do_sample: true,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
402
        max_new_tokens: default_max_new_tokens(),
403
        return_full_text: None,
404
        stop: Vec::new(),
405
        truncate: None,
406
        watermark: false,
OlivierDehaene's avatar
OlivierDehaene committed
407
        details: false,
408
        decoder_input_details: false,
409
        seed: None,
Nicolas Patry's avatar
Nicolas Patry committed
410
        top_n_tokens: None,
drbh's avatar
drbh committed
411
        grammar: None,
drbh's avatar
drbh committed
412
        adapter_id: None,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
413
414
415
    }
}

416
417
418
419
420
421
422
423
424
425
426
427
428
#[derive(Clone, Deserialize, Serialize, ToSchema, Debug)]
#[serde(try_from = "PromptDeserializer")]
pub struct Prompt(pub Vec<String>);

#[derive(Deserialize)]
#[serde(untagged)]
enum PromptDeserializer {
    Single(String),
    Multiple(Vec<String>),
}

impl TryFrom<PromptDeserializer> for Prompt {
    type Error = String;
429

430
    fn try_from(value: PromptDeserializer) -> Result<Self, Self::Error> {
431
        match value {
432
433
434
435
436
437
438
439
440
441
442
            PromptDeserializer::Single(s) => Ok(Prompt(vec![s])),
            PromptDeserializer::Multiple(v) => {
                if v.is_empty() {
                    Err(
                        "Empty array detected. Do not use an empty array for the prompt."
                            .to_string(),
                    )
                } else {
                    Ok(Prompt(v))
                }
            }
443
444
445
446
        }
    }
}

447
448
449
450
451
#[derive(Clone, Deserialize, Serialize, ToSchema, Debug)]
pub struct CompletionRequest {
    /// UNUSED
    #[schema(example = "mistralai/Mistral-7B-Instruct-v0.2")]
    /// ID of the model to use. See the model endpoint compatibility table for details on which models work with the Chat API.
452
    pub model: Option<String>,
453
454
455

    /// The prompt to generate completions for.
    #[schema(example = "What is Deep Learning?")]
456
    pub prompt: Prompt,
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493

    /// The maximum number of tokens that can be generated in the chat completion.
    #[serde(default)]
    #[schema(default = "32")]
    pub max_tokens: Option<u32>,

    /// What sampling temperature to use, between 0 and 2. Higher values like 0.8 will make the output more random, while
    /// lower values like 0.2 will make it more focused and deterministic. We generally recommend altering this or `top_p` but not both.
    #[serde(default)]
    #[schema(nullable = true, example = 1.0)]
    pub temperature: Option<f32>,

    /// An alternative to sampling with temperature, called nucleus sampling, where the model considers the results of the
    /// tokens with top_p probability mass. So 0.1 means only the tokens comprising the top 10% probability mass are considered.
    #[serde(default)]
    #[schema(nullable = true, example = 0.95)]
    pub top_p: Option<f32>,

    #[serde(default = "bool::default")]
    pub stream: bool,

    #[schema(nullable = true, example = 42)]
    pub seed: Option<u64>,

    /// The text to append to the prompt. This is useful for completing sentences or generating a paragraph of text.
    /// please see the completion_template field in the model's tokenizer_config.json file for completion template.
    #[serde(default)]
    pub suffix: Option<String>,

    #[serde(default)]
    pub repetition_penalty: Option<f32>,

    /// Number between -2.0 and 2.0. Positive values penalize new tokens based on their existing frequency in the text so far,
    /// decreasing the model's likelihood to repeat the same line verbatim.
    #[serde(default)]
    #[schema(example = "1.0")]
    pub frequency_penalty: Option<f32>,
494
495
496
497
498

    /// Up to 4 sequences where the API will stop generating further tokens.
    #[serde(default)]
    #[schema(nullable = true, example = "null")]
    pub stop: Option<Vec<String>>,
499
500
}

501
502
503
504
505
506
507
508
509
#[derive(Clone, Serialize, ToSchema)]
#[serde(tag = "object")]
enum Completion {
    #[serde(rename = "text_completion")]
    Chunk(Chunk),
    #[serde(rename = "text_completion")]
    Final(CompletionFinal),
}

510
#[derive(Clone, Deserialize, Serialize, ToSchema, Default)]
511
pub(crate) struct CompletionFinal {
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
    pub id: String,
    #[schema(example = "1706270835")]
    pub created: u64,
    #[schema(example = "mistralai/Mistral-7B-Instruct-v0.2")]
    pub model: String,
    pub system_fingerprint: String,
    pub choices: Vec<CompletionComplete>,
    pub usage: Usage,
}

#[derive(Clone, Deserialize, Serialize, ToSchema)]
pub(crate) struct CompletionComplete {
    pub index: u32,
    pub text: String,
    pub logprobs: Option<Vec<f32>>,
    pub finish_reason: String,
}

530
531
532
533
534
535
536
537
538
#[derive(Clone, Deserialize, Serialize, ToSchema)]
pub(crate) struct Chunk {
    pub id: String,
    pub created: u64,
    pub choices: Vec<CompletionComplete>,
    pub model: String,
    pub system_fingerprint: String,
}

539
#[derive(Clone, Deserialize, Serialize, ToSchema)]
540
541
pub(crate) struct ChatCompletion {
    pub id: String,
542
    #[schema(example = "1706270835")]
543
    pub created: u64,
544
    #[schema(example = "mistralai/Mistral-7B-Instruct-v0.2")]
545
546
547
548
549
550
    pub model: String,
    pub system_fingerprint: String,
    pub choices: Vec<ChatCompletionComplete>,
    pub usage: Usage,
}

551
#[derive(Clone, Deserialize, Serialize, ToSchema)]
552
553
pub(crate) struct ChatCompletionComplete {
    pub index: u32,
Nicolas Patry's avatar
Nicolas Patry committed
554
    pub message: OutputMessage,
555
    pub logprobs: Option<ChatCompletionLogprobs>,
556
557
558
    pub finish_reason: String,
}

559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
#[derive(Clone, Deserialize, Serialize, ToSchema)]
pub(crate) struct ChatCompletionLogprobs {
    content: Vec<ChatCompletionLogprob>,
}

impl From<(Token, Vec<Token>)> for ChatCompletionLogprobs {
    fn from(value: (Token, Vec<Token>)) -> Self {
        let (token, top_tokens) = value;

        Self {
            content: vec![ChatCompletionLogprob {
                token: token.text,
                logprob: token.logprob,
                top_logprobs: top_tokens
                    .into_iter()
                    .map(|t| ChatCompletionTopLogprob {
                        token: t.text,
                        logprob: t.logprob,
                    })
                    .collect(),
            }],
        }
    }
}

impl From<(Vec<Token>, Vec<Vec<Token>>)> for ChatCompletionLogprobs {
    fn from(value: (Vec<Token>, Vec<Vec<Token>>)) -> Self {
        let (tokens, top_tokens) = value;
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601

        // Create an iterator that produces None for top_tokens once it's exhausted
        let top_tokens_iter = top_tokens
            .into_iter()
            .map(Some)
            .chain(std::iter::repeat(None));

        let content = tokens
            .into_iter()
            .zip(top_tokens_iter)
            .map(|(t, top_t_option)| ChatCompletionLogprob {
                token: t.text,
                logprob: t.logprob,
                top_logprobs: match top_t_option {
                    Some(top_t) => top_t
602
603
604
605
606
607
                        .into_iter()
                        .map(|t| ChatCompletionTopLogprob {
                            token: t.text,
                            logprob: t.logprob,
                        })
                        .collect(),
608
609
610
611
612
613
                    None => vec![], // Handle the case where there are no top tokens
                },
            })
            .collect();

        Self { content }
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
    }
}

#[derive(Clone, Deserialize, Serialize, ToSchema)]
pub(crate) struct ChatCompletionLogprob {
    token: String,
    logprob: f32,
    top_logprobs: Vec<ChatCompletionTopLogprob>,
}

#[derive(Clone, Deserialize, Serialize, ToSchema)]
pub(crate) struct ChatCompletionTopLogprob {
    token: String,
    logprob: f32,
}

630
#[derive(Clone, Deserialize, Serialize, ToSchema, Default)]
631
632
633
634
635
636
pub(crate) struct Usage {
    pub prompt_tokens: u32,
    pub completion_tokens: u32,
    pub total_tokens: u32,
}

637
638
639
640
641
642
643
644
645
#[derive(Clone, Serialize, ToSchema)]
#[serde(tag = "object")]
enum CompletionType {
    #[serde(rename = "chat.completion.chunk")]
    ChatCompletionChunk(ChatCompletionChunk),
    #[serde(rename = "chat.completion")]
    ChatCompletion(ChatCompletion),
}

646
647
648
649
impl ChatCompletion {
    pub(crate) fn new(
        model: String,
        system_fingerprint: String,
drbh's avatar
drbh committed
650
        output: Option<String>,
651
652
653
        created: u64,
        details: Details,
        return_logprobs: bool,
654
        tool_calls: Option<Vec<ToolCall>>,
655
    ) -> Self {
Nicolas Patry's avatar
Nicolas Patry committed
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
        let message = match (output, tool_calls) {
            (Some(content), None) => OutputMessage::ChatMessage(TextMessage {
                role: "assistant".into(),
                content,
            }),
            (None, Some(tool_calls)) => OutputMessage::ToolCall(ToolCallMessage {
                role: "assistant".to_string(),
                tool_calls,
            }),
            (Some(output), Some(_)) => {
                warn!("Received both chat and tool call");
                OutputMessage::ChatMessage(TextMessage {
                    role: "assistant".into(),
                    content: output,
                })
            }
            (None, None) => {
                warn!("Didn't receive an answer");
                OutputMessage::ChatMessage(TextMessage {
                    role: "assistant".into(),
                    content: "".to_string(),
                })
            }
        };
680
681
682
683
684
685
686
        Self {
            id: String::new(),
            created,
            model,
            system_fingerprint,
            choices: vec![ChatCompletionComplete {
                index: 0,
Nicolas Patry's avatar
Nicolas Patry committed
687
                message,
688
                logprobs: return_logprobs
689
                    .then(|| ChatCompletionLogprobs::from((details.tokens, details.top_tokens))),
690
                finish_reason: details.finish_reason.format(true),
691
692
693
694
695
696
697
698
699
            }],
            usage: Usage {
                prompt_tokens: details.prefill.len() as u32,
                completion_tokens: details.generated_tokens,
                total_tokens: details.prefill.len() as u32 + details.generated_tokens,
            },
        }
    }
}
700
#[derive(Clone, Serialize, ToSchema)]
701
702
pub(crate) struct ChatCompletionChunk {
    pub id: String,
703
    #[schema(example = "1706270978")]
704
    pub created: u64,
705
    #[schema(example = "mistralai/Mistral-7B-Instruct-v0.2")]
706
707
708
    pub model: String,
    pub system_fingerprint: String,
    pub choices: Vec<ChatCompletionChoice>,
Nicolas Patry's avatar
Nicolas Patry committed
709
    pub usage: Option<Usage>,
710
711
}

712
#[derive(Clone, Serialize, ToSchema)]
713
714
715
pub(crate) struct ChatCompletionChoice {
    pub index: u32,
    pub delta: ChatCompletionDelta,
716
    pub logprobs: Option<ChatCompletionLogprobs>,
717
718
719
    pub finish_reason: Option<String>,
}

Nicolas Patry's avatar
Nicolas Patry committed
720
721
722
723
724
725
726
#[derive(Clone, Deserialize, ToSchema, Serialize, Debug, PartialEq)]
pub struct ToolCallDelta {
    #[schema(example = "assistant")]
    role: String,
    tool_calls: DeltaToolCall,
}

727
728
#[derive(Clone, Debug, Serialize, ToSchema)]
#[serde(untagged)]
Nicolas Patry's avatar
Nicolas Patry committed
729
730
731
enum ChatCompletionDelta {
    Chat(TextMessage),
    Tool(ToolCallDelta),
drbh's avatar
drbh committed
732
733
}

Nicolas Patry's avatar
Nicolas Patry committed
734
#[derive(Clone, Deserialize, Serialize, ToSchema, Debug, PartialEq)]
drbh's avatar
drbh committed
735
736
737
738
739
740
741
pub(crate) struct DeltaToolCall {
    pub index: u32,
    pub id: String,
    pub r#type: String,
    pub function: Function,
}

Nicolas Patry's avatar
Nicolas Patry committed
742
#[derive(Clone, Deserialize, Serialize, ToSchema, Debug, PartialEq)]
drbh's avatar
drbh committed
743
744
745
pub(crate) struct Function {
    pub name: Option<String>,
    pub arguments: String,
746
747
}

drbh's avatar
drbh committed
748
#[allow(clippy::too_many_arguments)]
749
750
751
752
impl ChatCompletionChunk {
    pub(crate) fn new(
        model: String,
        system_fingerprint: String,
drbh's avatar
drbh committed
753
754
        delta: Option<String>,
        tool_calls: Option<Vec<String>>,
755
        created: u64,
756
        logprobs: Option<ChatCompletionLogprobs>,
757
        finish_reason: Option<String>,
Nicolas Patry's avatar
Nicolas Patry committed
758
        usage: Option<Usage>,
759
    ) -> Self {
760
        let delta = match (delta, tool_calls) {
Nicolas Patry's avatar
Nicolas Patry committed
761
762
763
764
765
766
767
            (Some(delta), _) => ChatCompletionDelta::Chat(TextMessage {
                role: "assistant".to_string(),
                content: delta,
            }),
            (None, Some(tool_calls)) => ChatCompletionDelta::Tool(ToolCallDelta {
                role: "assistant".to_string(),
                tool_calls: DeltaToolCall {
768
769
770
771
772
773
774
                    index: 0,
                    id: String::new(),
                    r#type: "function".to_string(),
                    function: Function {
                        name: None,
                        arguments: tool_calls[0].to_string(),
                    },
Nicolas Patry's avatar
Nicolas Patry committed
775
776
777
778
779
780
                },
            }),
            (None, None) => ChatCompletionDelta::Chat(TextMessage {
                role: "assistant".to_string(),
                content: "".to_string(),
            }),
781
        };
782
783
784
785
786
787
        Self {
            id: String::new(),
            created,
            model,
            system_fingerprint,
            choices: vec![ChatCompletionChoice {
788
                index: 0,
789
                delta,
790
791
792
                logprobs,
                finish_reason,
            }],
Nicolas Patry's avatar
Nicolas Patry committed
793
            usage,
794
795
796
797
798
        }
    }
}

#[derive(Clone, Deserialize, ToSchema, Serialize)]
Nicolas Patry's avatar
Nicolas Patry committed
799
#[cfg_attr(test, derive(Debug, PartialEq, Default))]
800
pub(crate) struct ChatRequest {
801
    #[schema(example = "mistralai/Mistral-7B-Instruct-v0.2")]
drbh's avatar
drbh committed
802
    /// [UNUSED] ID of the model to use. See the model endpoint compatibility table for details on which models work with the Chat API.
803
    pub model: Option<String>,
drbh's avatar
drbh committed
804

805
    /// A list of messages comprising the conversation so far.
drbh's avatar
drbh committed
806
    #[schema(example = "[{\"role\": \"user\", \"content\": \"What is Deep Learning?\"}]")]
807
808
809
810
811
    pub messages: Vec<Message>,

    /// Number between -2.0 and 2.0. Positive values penalize new tokens based on their existing frequency in the text so far,
    /// decreasing the model's likelihood to repeat the same line verbatim.
    #[serde(default)]
812
    #[schema(example = "1.0")]
813
814
815
816
817
818
819
820
821
822
823
824
825
826
    pub frequency_penalty: Option<f32>,

    /// UNUSED
    /// Modify the likelihood of specified tokens appearing in the completion. Accepts a JSON object that maps tokens
    /// (specified by their token ID in the tokenizer) to an associated bias value from -100 to 100. Mathematically,
    /// the bias is added to the logits generated by the model prior to sampling. The exact effect will vary per model,
    /// but values between -1 and 1 should decrease or increase likelihood of selection; values like -100 or 100 should
    /// result in a ban or exclusive selection of the relevant token.
    #[serde(default)]
    pub logit_bias: Option<Vec<f32>>,

    /// Whether to return log probabilities of the output tokens or not. If true, returns the log probabilities of each
    /// output token returned in the content of message.
    #[serde(default)]
827
    #[schema(example = "false")]
828
829
830
831
832
    pub logprobs: Option<bool>,

    /// An integer between 0 and 5 specifying the number of most likely tokens to return at each token position, each with
    /// an associated log probability. logprobs must be set to true if this parameter is used.
    #[serde(default)]
833
    #[schema(example = "5")]
834
835
836
837
    pub top_logprobs: Option<u32>,

    /// The maximum number of tokens that can be generated in the chat completion.
    #[serde(default)]
838
    #[schema(example = "32")]
839
840
841
842
843
844
    pub max_tokens: Option<u32>,

    /// UNUSED
    /// How many chat completion choices to generate for each input message. Note that you will be charged based on the
    /// number of generated tokens across all of the choices. Keep n as 1 to minimize costs.
    #[serde(default)]
845
    #[schema(nullable = true, example = "2")]
846
847
848
849
850
    pub n: Option<u32>,

    /// Number between -2.0 and 2.0. Positive values penalize new tokens based on whether they appear in the text so far,
    /// increasing the model's likelihood to talk about new topics
    #[serde(default)]
851
    #[schema(nullable = true, example = 0.1)]
852
853
    pub presence_penalty: Option<f32>,

854
855
856
857
858
    /// Up to 4 sequences where the API will stop generating further tokens.
    #[serde(default)]
    #[schema(nullable = true, example = "null")]
    pub stop: Option<Vec<String>>,

859
860
861
862
863
    #[serde(default = "bool::default")]
    pub stream: bool,

    #[schema(nullable = true, example = 42)]
    pub seed: Option<u64>,
864
865
866
867
868
869

    /// What sampling temperature to use, between 0 and 2. Higher values like 0.8 will make the output more random, while
    /// lower values like 0.2 will make it more focused and deterministic.
    ///
    /// We generally recommend altering this or `top_p` but not both.
    #[serde(default)]
870
    #[schema(nullable = true, example = 1.0)]
871
872
873
874
875
    pub temperature: Option<f32>,

    /// An alternative to sampling with temperature, called nucleus sampling, where the model considers the results of the
    /// tokens with top_p probability mass. So 0.1 means only the tokens comprising the top 10% probability mass are considered.
    #[serde(default)]
876
    #[schema(nullable = true, example = 0.95)]
877
    pub top_p: Option<f32>,
drbh's avatar
drbh committed
878
879
880
881
882
883
884
885

    /// A list of tools the model may call. Currently, only functions are supported as a tool. Use this to provide a list of
    /// functions the model may generate JSON inputs for.
    #[serde(default)]
    #[schema(nullable = true, example = "null")]
    pub tools: Option<Vec<Tool>>,

    /// A prompt to be appended before the tools
drbh's avatar
drbh committed
886
    #[serde(default)]
drbh's avatar
drbh committed
887
888
    #[schema(
        nullable = true,
drbh's avatar
drbh committed
889
        example = "Given the functions available, please respond with a JSON for a function call with its proper arguments that best answers the given prompt. Respond in the format {name: function name, parameters: dictionary of argument name and its value}.Do not use variables."
drbh's avatar
drbh committed
890
891
892
893
894
895
    )]
    pub tool_prompt: Option<String>,

    /// A specific tool to use. If not provided, the model will default to use any of the tools provided in the tools parameter.
    #[serde(default)]
    #[schema(nullable = true, example = "null")]
drbh's avatar
drbh committed
896
    pub tool_choice: ToolChoice,
drbh's avatar
drbh committed
897
898
899
900
901
902
903

    /// Response format constraints for the generation.
    ///
    /// NOTE: A request can use `response_format` OR `tools` but not both.
    #[serde(default)]
    #[schema(nullable = true, default = "null", example = "null")]
    pub response_format: Option<GrammarType>,
904
905
906
907
908

    /// A guideline to be used in the chat_template
    #[serde(default)]
    #[schema(nullable = true, default = "null", example = "null")]
    pub guideline: Option<String>,
Nicolas Patry's avatar
Nicolas Patry committed
909
910
911
912
913
914
915

    /// Options for streaming response. Only set this when you set stream: true.
    #[serde(default)]
    #[schema(nullable = true, example = "null")]
    pub stream_options: Option<StreamOptions>,
}

Nicolas Patry's avatar
Nicolas Patry committed
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
impl ChatRequest {
    fn try_into_generate(self, infer: &Infer) -> Result<(GenerateRequest, bool), InferError> {
        let ChatRequest {
            model,
            max_tokens,
            messages,
            seed,
            stop,
            stream,
            tools,
            tool_choice,
            tool_prompt,
            temperature,
            response_format,
            guideline,
            presence_penalty,
            frequency_penalty,
            top_p,
            top_logprobs,
            ..
        } = self;

        let repetition_penalty = presence_penalty.map(|x| x + 2.0);
        let max_new_tokens = max_tokens.or(Some(100));
        let tool_prompt = tool_prompt
            .filter(|s| !s.is_empty())
            .unwrap_or_else(default_tool_prompt);
        let stop = stop.unwrap_or_default();
        // enable greedy only when temperature is 0
        let (do_sample, temperature) = match temperature {
            Some(temperature) if temperature == 0.0 => (false, None),
            other => (true, other),
        };
        let (inputs, grammar, using_tools) = prepare_chat_input(
            infer,
            response_format,
            tools,
            tool_choice,
            &tool_prompt,
            guideline,
            messages,
        )?;

        Ok((
            GenerateRequest {
                inputs: inputs.to_string(),
                add_special_tokens: false,
                parameters: GenerateParameters {
                    best_of: None,
                    temperature,
                    repetition_penalty,
                    frequency_penalty,
                    top_k: None,
                    top_p,
                    typical_p: None,
                    do_sample,
                    max_new_tokens,
                    return_full_text: None,
                    stop,
                    truncate: None,
                    watermark: false,
                    details: true,
                    decoder_input_details: !stream,
                    seed,
                    top_n_tokens: top_logprobs,
                    grammar,
                    adapter_id: model.filter(|m| *m != "tgi").map(String::from),
                },
            },
            using_tools,
        ))
    }
}

Nicolas Patry's avatar
Nicolas Patry committed
990
#[derive(Clone, Deserialize, ToSchema, Serialize)]
Nicolas Patry's avatar
Nicolas Patry committed
991
#[cfg_attr(test, derive(Debug, PartialEq))]
Nicolas Patry's avatar
Nicolas Patry committed
992
993
994
995
struct StreamOptions {
    /// If set, an additional chunk will be streamed before the data: [DONE] message. The usage field on this chunk shows the token usage statistics for the entire request, and the choices field will always be an empty array. All other chunks will also include a usage field, but with a null value.
    #[schema(example = "true")]
    include_usage: bool,
drbh's avatar
drbh committed
996
997
}

drbh's avatar
drbh committed
998
999
pub fn default_tool_prompt() -> String {
    "\nGiven the functions available, please respond with a JSON for a function call with its proper arguments that best answers the given prompt. Respond in the format {name: function name, parameters: dictionary of argument name and its value}.Do not use variables.\n".to_string()
drbh's avatar
drbh committed
1000
}
1001
1002

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize, ToSchema)]
1003
1004
#[schema(example = "auto")]
/// Controls which (if any) tool is called by the model.
1005
pub enum ToolType {
1006
1007
    /// Means the model can pick between generating a message or calling one or more tools.
    #[schema(rename = "auto")]
drbh's avatar
drbh committed
1008
    OneOf,
1009
1010
    /// Means the model will not call any tool and instead generates a message.
    #[schema(rename = "none")]
drbh's avatar
drbh committed
1011
    NoTool,
1012
1013
1014
    /// Forces the model to call a specific tool.
    #[schema(rename = "function")]
    Function(FunctionName),
drbh's avatar
drbh committed
1015
1016
}

1017
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema)]
1018
1019
1020
1021
pub struct FunctionName {
    pub name: String,
}

drbh's avatar
drbh committed
1022
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default, ToSchema)]
1023
1024
#[serde(from = "ToolTypeDeserializer")]
pub struct ToolChoice(pub Option<ToolType>);
drbh's avatar
drbh committed
1025

1026
1027
1028
#[derive(Deserialize)]
#[serde(untagged)]
enum ToolTypeDeserializer {
1029
    Null,
drbh's avatar
drbh committed
1030
1031
    String(String),
    ToolType(ToolType),
1032
}
drbh's avatar
drbh committed
1033

1034
1035
impl From<ToolTypeDeserializer> for ToolChoice {
    fn from(value: ToolTypeDeserializer) -> Self {
drbh's avatar
drbh committed
1036
        match value {
1037
            ToolTypeDeserializer::Null => ToolChoice(None),
drbh's avatar
drbh committed
1038
1039
1040
            ToolTypeDeserializer::String(s) => match s.as_str() {
                "none" => ToolChoice(Some(ToolType::NoTool)),
                "auto" => ToolChoice(Some(ToolType::OneOf)),
1041
                _ => ToolChoice(Some(ToolType::Function(FunctionName { name: s }))),
drbh's avatar
drbh committed
1042
            },
drbh's avatar
drbh committed
1043
            ToolTypeDeserializer::ToolType(tool_type) => ToolChoice(Some(tool_type)),
drbh's avatar
drbh committed
1044
1045
1046
1047
        }
    }
}

1048
#[derive(Debug, Deserialize, Serialize, ToSchema, PartialEq)]
drbh's avatar
drbh committed
1049
pub struct JsonSchemaTool {
drbh's avatar
drbh committed
1050
1051
1052
1053
1054
    #[serde(flatten)]
    functions_map: FunctionsMap,
    properties: Properties,
}

1055
#[derive(Debug, Serialize, Deserialize, PartialEq)]
drbh's avatar
drbh committed
1056
1057
1058
1059
1060
struct FunctionsMap {
    #[serde(rename = "$functions")]
    functions: std::collections::HashMap<String, serde_json::Value>,
}

1061
#[derive(Debug, Serialize, Deserialize, PartialEq)]
drbh's avatar
drbh committed
1062
1063
1064
1065
1066
struct FunctionRef {
    #[serde(rename = "$ref")]
    ref_path: String,
}

1067
#[derive(Debug, Serialize, Deserialize, PartialEq)]
drbh's avatar
drbh committed
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
struct Properties {
    #[serde(serialize_with = "serialize_function")]
    function: Vec<FunctionRef>,
}

fn serialize_function<S>(functions: &Vec<FunctionRef>, serializer: S) -> Result<S::Ok, S::Error>
where
    S: serde::Serializer,
{
    use serde::ser::SerializeStruct;
    let mut state = serializer.serialize_struct("Function", 1)?;
    state.serialize_field("anyOf", functions)?;
    state.end()
}

Nicolas Patry's avatar
Nicolas Patry committed
1083
#[derive(Clone, Debug, Deserialize, Serialize, ToSchema, Default, PartialEq)]
drbh's avatar
drbh committed
1084
1085
1086
1087
pub(crate) struct FunctionDefinition {
    #[serde(default)]
    pub description: Option<String>,
    pub name: String,
1088
1089
    #[serde(alias = "parameters")]
    pub arguments: serde_json::Value,
drbh's avatar
drbh committed
1090
1091
1092
}

#[derive(Clone, Debug, Deserialize, Serialize, ToSchema)]
Nicolas Patry's avatar
Nicolas Patry committed
1093
#[cfg_attr(test, derive(PartialEq))]
drbh's avatar
drbh committed
1094
1095
1096
1097
1098
1099
pub(crate) struct Tool {
    // The type of the tool. Currently, only 'function' is supported.
    #[schema(example = "function")]
    pub r#type: String,
    // Grab the tool as generic JSON for debugging purposes.
    pub function: FunctionDefinition,
1100
1101
}

1102
#[derive(Clone, Serialize, Deserialize, Default)]
1103
pub(crate) struct ChatTemplateInputs<'a> {
Nicolas Patry's avatar
Nicolas Patry committed
1104
    messages: Vec<TextMessage>,
1105
1106
    bos_token: Option<&'a str>,
    eos_token: Option<&'a str>,
1107
    add_generation_prompt: bool,
drbh's avatar
drbh committed
1108
    tools: Option<Vec<Tool>>,
1109
    guideline: Option<&'a str>,
1110
1111
}

Nicolas Patry's avatar
Nicolas Patry committed
1112
#[derive(Clone, Deserialize, Serialize, ToSchema, Default, Debug, PartialEq)]
drbh's avatar
drbh committed
1113
pub(crate) struct ToolCall {
1114
    pub id: String,
drbh's avatar
drbh committed
1115
1116
1117
1118
    pub r#type: String,
    pub function: FunctionDefinition,
}

Nicolas Patry's avatar
Nicolas Patry committed
1119
#[derive(Clone, Deserialize, ToSchema, Serialize, Debug, PartialEq)]
1120
pub struct Url {
Nicolas Patry's avatar
Nicolas Patry committed
1121
    url: String,
drbh's avatar
drbh committed
1122
1123
}

Nicolas Patry's avatar
Nicolas Patry committed
1124
1125
1126
#[derive(Clone, Deserialize, ToSchema, Serialize, Debug, PartialEq)]
#[serde(tag = "type")]
#[serde(rename_all = "snake_case")]
1127
1128
1129
pub enum MessageChunk {
    Text { text: String },
    ImageUrl { image_url: Url },
Nicolas Patry's avatar
Nicolas Patry committed
1130
1131
1132
1133
1134
1135
1136
}

#[derive(Clone, Deserialize, ToSchema, Serialize, Debug, PartialEq)]
pub struct Message {
    #[schema(example = "user")]
    role: String,
    #[schema(example = "My name is David and I")]
1137
    pub content: MessageContent,
drbh's avatar
drbh committed
1138
    #[serde(default, skip_serializing_if = "Option::is_none")]
Nicolas Patry's avatar
Nicolas Patry committed
1139
1140
    #[schema(example = "\"David\"")]
    name: Option<String>,
drbh's avatar
drbh committed
1141
1142
}

1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
#[derive(Clone, Deserialize, Serialize, ToSchema, Debug, PartialEq)]
#[serde(untagged)]
pub enum MessageContent {
    SingleText(String),
    MultipleChunks(Vec<MessageChunk>),
}

// Pushing a chunk to a single text message will convert it to a multiple chunks message
impl MessageContent {
    pub fn push(&mut self, chunk: MessageChunk) {
        match self {
            MessageContent::SingleText(text) => {
drbh's avatar
drbh committed
1155
1156
1157
1158
                *self = MessageContent::MultipleChunks(vec![
                    MessageChunk::Text { text: text.clone() },
                    chunk,
                ]);
Nicolas Patry's avatar
Nicolas Patry committed
1159
            }
1160
1161
1162
1163
            MessageContent::MultipleChunks(chunks) => {
                chunks.push(chunk);
            }
        }
drbh's avatar
drbh committed
1164
1165
1166
    }
}

Nicolas Patry's avatar
Nicolas Patry committed
1167
1168
#[derive(Clone, Deserialize, ToSchema, Serialize, Debug, PartialEq)]
pub struct TextMessage {
1169
1170
1171
    #[schema(example = "user")]
    pub role: String,
    #[schema(example = "My name is David and I")]
Nicolas Patry's avatar
Nicolas Patry committed
1172
1173
1174
1175
1176
1177
1178
    pub content: String,
}

impl From<Message> for TextMessage {
    fn from(value: Message) -> Self {
        TextMessage {
            role: value.role,
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
            content: match value.content {
                MessageContent::SingleText(text) => text,
                MessageContent::MultipleChunks(chunks) => chunks
                    .into_iter()
                    .map(|chunk| match chunk {
                        MessageChunk::Text { text } => text,
                        MessageChunk::ImageUrl { image_url } => format!("![]({})", image_url.url),
                    })
                    .collect::<Vec<_>>()
                    .join(""),
            },
Nicolas Patry's avatar
Nicolas Patry committed
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
        }
    }
}

#[derive(Clone, Deserialize, ToSchema, Serialize, Debug, PartialEq)]
pub struct ToolCallMessage {
    #[schema(example = "assistant")]
    role: String,
    tool_calls: Vec<ToolCall>,
}

#[derive(Clone, Deserialize, ToSchema, Serialize, Debug, PartialEq)]
#[serde(untagged)]
pub(crate) enum OutputMessage {
    ChatMessage(TextMessage),
    ToolCall(ToolCallMessage),
1206
1207
}

1208
#[derive(Clone, Debug, Deserialize, ToSchema)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1209
pub(crate) struct GenerateRequest {
1210
    #[schema(example = "My name is Olivier and I")]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1211
1212
1213
    pub inputs: String,
    #[serde(default = "default_parameters")]
    pub parameters: GenerateParameters,
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223

    /// This is used internally because some requests
    /// already contain the templated input therefore
    /// we shouldn't add the special tokens.
    #[serde(default = "default_true", skip)]
    pub add_special_tokens: bool,
}

fn default_true() -> bool {
    true
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1224
1225
}

1226
1227
1228
1229
1230
1231
1232
#[derive(Clone, Debug, Deserialize, ToSchema)]
pub(crate) struct CompatGenerateRequest {
    #[schema(example = "My name is Olivier and I")]
    pub inputs: String,
    #[serde(default = "default_parameters")]
    pub parameters: GenerateParameters,
    #[serde(default)]
OlivierDehaene's avatar
OlivierDehaene committed
1233
    #[schema(default = "false")]
1234
1235
1236
1237
1238
1239
1240
    pub stream: bool,
}

impl From<CompatGenerateRequest> for GenerateRequest {
    fn from(req: CompatGenerateRequest) -> Self {
        Self {
            inputs: req.inputs,
1241
            add_special_tokens: true,
1242
1243
1244
1245
1246
            parameters: req.parameters,
        }
    }
}

1247
1248
1249
#[derive(Debug, Serialize, ToSchema)]
pub struct PrefillToken {
    #[schema(example = 0)]
Nicolas Patry's avatar
Nicolas Patry committed
1250
    pub id: u32,
1251
    #[schema(example = "test")]
Nicolas Patry's avatar
Nicolas Patry committed
1252
    pub text: String,
1253
    #[schema(nullable = true, example = - 0.34)]
Nicolas Patry's avatar
Nicolas Patry committed
1254
    pub logprob: f32,
1255
1256
}

1257
#[derive(Debug, Serialize, ToSchema, Clone)]
1258
1259
pub struct Token {
    #[schema(example = 0)]
Nicolas Patry's avatar
Nicolas Patry committed
1260
    pub id: u32,
1261
    #[schema(example = "test")]
Nicolas Patry's avatar
Nicolas Patry committed
1262
    pub text: String,
1263
    #[schema(nullable = true, example = - 0.34)]
Nicolas Patry's avatar
Nicolas Patry committed
1264
    pub logprob: f32,
1265
    #[schema(example = "false")]
Nicolas Patry's avatar
Nicolas Patry committed
1266
    pub special: bool,
1267
1268
}

1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
#[derive(Debug, Serialize, ToSchema)]
pub struct SimpleToken {
    #[schema(example = 0)]
    id: u32,
    #[schema(example = "test")]
    text: String,
    #[schema(example = 0)]
    start: usize,
    #[schema(example = 2)]
    stop: usize,
}

OlivierDehaene's avatar
OlivierDehaene committed
1281
#[derive(Debug, Serialize, ToSchema)]
1282
#[serde(rename_all(serialize = "snake_case"))]
1283
#[schema(example = "Length")]
Nicolas Patry's avatar
Nicolas Patry committed
1284
pub enum FinishReason {
1285
1286
1287
1288
1289
1290
1291
1292
    #[schema(rename = "length")]
    Length,
    #[serde(rename = "eos_token")]
    #[schema(rename = "eos_token")]
    EndOfSequenceToken,
    #[schema(rename = "stop_sequence")]
    StopSequence,
}
1293

1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
impl std::fmt::Display for FinishReason {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            FinishReason::Length => write!(f, "length"),
            FinishReason::EndOfSequenceToken => write!(f, "eos_token"),
            FinishReason::StopSequence => write!(f, "stop_sequence"),
        }
    }
}

1304
1305
1306
1307
1308
1309
1310
1311
1312
impl FinishReason {
    pub fn format(&self, use_stop: bool) -> String {
        match self {
            FinishReason::EndOfSequenceToken if use_stop => "stop".to_string(),
            _ => self.to_string(),
        }
    }
}

1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
#[derive(Serialize, ToSchema)]
pub(crate) struct BestOfSequence {
    #[schema(example = "test")]
    pub generated_text: String,
    #[schema(example = "length")]
    pub finish_reason: FinishReason,
    #[schema(example = 1)]
    pub generated_tokens: u32,
    #[schema(nullable = true, example = 42)]
    pub seed: Option<u64>,
    pub prefill: Vec<PrefillToken>,
    pub tokens: Vec<Token>,
Nicolas Patry's avatar
Nicolas Patry committed
1325
1326
    #[serde(skip_serializing_if = "Vec::is_empty")]
    pub top_tokens: Vec<Vec<Token>>,
1327
1328
}

1329
#[derive(Serialize, ToSchema)]
OlivierDehaene's avatar
OlivierDehaene committed
1330
pub(crate) struct Details {
1331
1332
1333
    #[schema(example = "length")]
    pub finish_reason: FinishReason,
    #[schema(example = 1)]
OlivierDehaene's avatar
OlivierDehaene committed
1334
    pub generated_tokens: u32,
1335
    #[schema(nullable = true, example = 42)]
1336
    pub seed: Option<u64>,
1337
1338
    pub prefill: Vec<PrefillToken>,
    pub tokens: Vec<Token>,
1339
1340
    #[serde(skip_serializing_if = "Option::is_none")]
    pub best_of_sequences: Option<Vec<BestOfSequence>>,
Nicolas Patry's avatar
Nicolas Patry committed
1341
1342
    #[serde(skip_serializing_if = "Vec::is_empty")]
    pub top_tokens: Vec<Vec<Token>>,
OlivierDehaene's avatar
OlivierDehaene committed
1343
1344
}

1345
#[derive(Serialize, ToSchema)]
1346
pub(crate) struct GenerateResponse {
1347
    #[schema(example = "test")]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1348
    pub generated_text: String,
OlivierDehaene's avatar
OlivierDehaene committed
1349
1350
    #[serde(skip_serializing_if = "Option::is_none")]
    pub details: Option<Details>,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1351
}
1352

1353
1354
1355
1356
1357
1358
#[derive(Serialize, ToSchema)]
pub(crate) struct ChatTokenizeResponse {
    pub(crate) tokenize_response: TokenizeResponse,
    pub(crate) templated_text: String,
}

1359
1360
1361
1362
#[derive(Serialize, ToSchema)]
#[serde(transparent)]
pub(crate) struct TokenizeResponse(Vec<SimpleToken>);

1363
1364
1365
1366
1367
1368
#[derive(Serialize, ToSchema)]
pub(crate) struct StreamDetails {
    #[schema(example = "length")]
    pub finish_reason: FinishReason,
    #[schema(example = 1)]
    pub generated_tokens: u32,
1369
    #[schema(nullable = true, example = 42)]
1370
    pub seed: Option<u64>,
1371
1372
    #[schema(example = 1)]
    pub input_length: u32,
1373
1374
1375
}

#[derive(Serialize, ToSchema)]
1376
pub(crate) struct StreamResponse {
1377
    pub index: u32,
1378
    pub token: Token,
Nicolas Patry's avatar
Nicolas Patry committed
1379
1380
    #[serde(skip_serializing_if = "Vec::is_empty")]
    pub top_tokens: Vec<Token>,
1381
    #[schema(nullable = true, default = "null", example = "test")]
1382
    pub generated_text: Option<String>,
1383
1384
    #[schema(nullable = true, default = "null")]
    pub details: Option<StreamDetails>,
1385
1386
}

1387
#[derive(Serialize, ToSchema)]
1388
1389
pub(crate) struct ErrorResponse {
    pub error: String,
1390
    pub error_type: String,
1391
}
1392

drbh's avatar
drbh committed
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
#[derive(Serialize, Deserialize, ToSchema)]
pub(crate) struct ModelInfo {
    #[schema(example = "gpt2")]
    pub id: String,
    #[schema(example = "model")]
    pub object: String,
    #[schema(example = 1686935002)]
    pub created: u64,
    #[schema(example = "openai")]
    pub owned_by: String,
}

#[derive(Serialize, Deserialize, ToSchema)]
pub(crate) struct ModelsInfo {
    #[schema(example = "list")]
    pub object: String,
    pub data: Vec<ModelInfo>,
}

impl Default for ModelsInfo {
    fn default() -> Self {
        ModelsInfo {
            object: "list".to_string(),
            data: Vec::new(),
        }
    }
}

1421
#[cfg(test)]
1422
mod tests {
1423
    use super::*;
Nicolas Patry's avatar
Nicolas Patry committed
1424
    use serde_json::json;
1425

1426
    pub(crate) fn get_tokenizer() -> Tokenizer {
1427
1428
1429
        let api = hf_hub::api::sync::Api::new().unwrap();
        let repo = api.model("gpt2".to_string());
        let filename = repo.get("tokenizer.json").unwrap();
1430
        Tokenizer::Rust(tokenizers::Tokenizer::from_file(filename).unwrap())
1431
    }
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445

    #[test]
    fn test_hub_nested_tokens_tokenizer_config() {
        // this is a subset of the tokenizer.json file
        // in this case we expect the tokens to be encoded as simple strings
        let json_content = r#"{
            "chat_template": "test",
            "bos_token": "<|begin▁of▁sentence|>",
            "eos_token": "<|end▁of▁sentence|>"
        }"#;

        let config: HubTokenizerConfig = serde_json::from_str(json_content).unwrap();

        // check that we successfully parsed the tokens
1446
1447
1448
1449
        assert_eq!(
            config.chat_template,
            Some(ChatTemplateVersions::Single("test".to_string()))
        );
1450
1451
        assert_eq!(
            config.bos_token,
1452
1453
1454
1455
1456
1457
1458
1459
1460
            Some(TokenizerConfigToken::String(
                "<|begin▁of▁sentence|>".to_string()
            ))
        );
        assert_eq!(
            config.eos_token,
            Some(TokenizerConfigToken::String(
                "<|end▁of▁sentence|>".to_string()
            ))
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
        );

        // in this case we expect the tokens to be encoded as structured tokens
        // we want the content of the structured token
        let json_content = r#"{
            "chat_template": "test",
            "bos_token": {
              "__type": "AddedToken",
              "content": "<|begin▁of▁sentence|>",
              "lstrip": false,
              "normalized": true,
              "rstrip": false,
              "single_word": false
            },
            "eos_token": {
              "__type": "AddedToken",
              "content": "<|end▁of▁sentence|>",
              "lstrip": false,
              "normalized": true,
              "rstrip": false,
              "single_word": false
            }
        }"#;

        let config: HubTokenizerConfig = serde_json::from_str(json_content).unwrap();

        // check that we successfully parsed the tokens
1488
1489
1490
1491
        assert_eq!(
            config.chat_template,
            Some(ChatTemplateVersions::Single("test".to_string()))
        );
1492
1493
        assert_eq!(
            config.bos_token,
1494
1495
1496
1497
1498
1499
1500
1501
1502
            Some(TokenizerConfigToken::Object {
                content: "<|begin▁of▁sentence|>".to_string()
            })
        );
        assert_eq!(
            config.eos_token,
            Some(TokenizerConfigToken::Object {
                content: "<|end▁of▁sentence|>".to_string()
            })
1503
1504
        );
    }
Nicolas Patry's avatar
Nicolas Patry committed
1505
1506
1507

    #[test]
    fn test_chat_simple_string() {
Nicolas Patry's avatar
Nicolas Patry committed
1508
        let json = json!({
Nicolas Patry's avatar
Nicolas Patry committed
1509
            "model": "",
Nicolas Patry's avatar
Nicolas Patry committed
1510
1511
            "messages": [{
                "role": "user",
Nicolas Patry's avatar
Nicolas Patry committed
1512
                "content": "What is Deep Learning?"
Nicolas Patry's avatar
Nicolas Patry committed
1513
            }]
Nicolas Patry's avatar
Nicolas Patry committed
1514
1515
1516
1517
1518
1519
1520
        });
        let request: ChatRequest = serde_json::from_str(json.to_string().as_str()).unwrap();

        assert_eq!(
            request.messages[0],
            Message {
                role: "user".to_string(),
1521
                content: MessageContent::SingleText("What is Deep Learning?".to_string()),
Nicolas Patry's avatar
Nicolas Patry committed
1522
1523
1524
1525
1526
                name: None
            }
        );
    }

drbh's avatar
drbh committed
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
    #[test]
    fn test_message_content_append() {
        let mut content = MessageContent::SingleText("Initial text".to_string());
        let chunk = MessageChunk::Text {
            text: "Additional text".to_string(),
        };

        content.push(chunk);

        match content {
            MessageContent::MultipleChunks(chunks) => {
                assert_eq!(chunks.len(), 2);
                assert_eq!(
                    chunks[0],
                    MessageChunk::Text {
                        text: "Initial text".to_string()
                    }
                );
                assert_eq!(
                    chunks[1],
                    MessageChunk::Text {
                        text: "Additional text".to_string()
                    }
                );
            }
            _ => panic!("Expected MultipleChunks, but got a different variant"),
        }
    }

Nicolas Patry's avatar
Nicolas Patry committed
1556
1557
    #[test]
    fn test_chat_request() {
Nicolas Patry's avatar
Nicolas Patry committed
1558
        let json = json!({
Nicolas Patry's avatar
Nicolas Patry committed
1559
            "model": "",
Nicolas Patry's avatar
Nicolas Patry committed
1560
1561
            "messages": [{
                "role": "user",
Nicolas Patry's avatar
Nicolas Patry committed
1562
1563
                "content": [
                    {"type": "text", "text": "Whats in this image?"},
Nicolas Patry's avatar
Nicolas Patry committed
1564
                    {"type": "image_url", "image_url": {"url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/rabbit.png"}},
Nicolas Patry's avatar
Nicolas Patry committed
1565
                ]
Nicolas Patry's avatar
Nicolas Patry committed
1566
            }]
Nicolas Patry's avatar
Nicolas Patry committed
1567
1568
1569
1570
1571
1572
1573
        });
        let request: ChatRequest = serde_json::from_str(json.to_string().as_str()).unwrap();

        assert_eq!(
            request.messages[0],
            Message{
                role: "user".to_string(),
1574
1575
1576
1577
                content: MessageContent::MultipleChunks(vec![
                    MessageChunk::Text { text: "Whats in this image?".to_string() },
                    MessageChunk::ImageUrl { image_url: Url { url: "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/rabbit.png".to_string() }},
                ]),
Nicolas Patry's avatar
Nicolas Patry committed
1578
1579
1580
1581
                name: None
            }
        );
    }
Nicolas Patry's avatar
Nicolas Patry committed
1582
1583
1584
1585
1586

    #[test]
    fn text_message_convert() {
        let message = Message{
                role: "user".to_string(),
1587
1588
1589
1590
                content: MessageContent::MultipleChunks(vec![
                    MessageChunk::Text { text: "Whats in this image?".to_string() },
                    MessageChunk::ImageUrl { image_url: Url { url: "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/rabbit.png".to_string() } }
                ]),
Nicolas Patry's avatar
Nicolas Patry committed
1591
1592
1593
1594
1595
                name: None
            };
        let textmsg: TextMessage = message.into();
        assert_eq!(textmsg.content, "Whats in this image?![](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/rabbit.png)");
    }
Nicolas Patry's avatar
Nicolas Patry committed
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616

    #[test]
    fn test_chat_stream_options() {
        let json = json!({
            "model": "",
            "stream_options": {"include_usage": true},
            "messages": [{
                "role": "user",
                "content": "Hello"
            }]
        });
        let request: ChatRequest = serde_json::from_str(json.to_string().as_str()).unwrap();

        assert!(matches!(
            request.stream_options,
            Some(StreamOptions {
                include_usage: true
            })
        ));
    }

Nicolas Patry's avatar
Nicolas Patry committed
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
    #[test]
    fn openai_output() {
        let message = OutputMessage::ChatMessage(TextMessage {
            role: "assistant".to_string(),
            content: "This is the answer".to_string(),
        });
        let serialized = serde_json::to_string(&message).unwrap();
        assert_eq!(
            serialized,
            r#"{"role":"assistant","content":"This is the answer"}"#
        );

        let message = OutputMessage::ToolCall(ToolCallMessage {
            role: "assistant".to_string(),
            tool_calls: vec![ToolCall {
                id: "0".to_string(),
                r#type: "function".to_string(),
                function: FunctionDefinition {
                    description: None,
                    name: "myfn".to_string(),
                    arguments: json!({
                        "format": "csv"
                    }),
                },
            }],
        });
        let serialized = serde_json::to_string(&message).unwrap();
        assert_eq!(
            serialized,
            r#"{"role":"assistant","tool_calls":[{"id":"0","type":"function","function":{"description":null,"name":"myfn","arguments":{"format":"csv"}}}]}"#
        );
    }
1649
}