lib.rs 45.4 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
12
pub mod usage_stats;

Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
13
use serde::{Deserialize, Serialize};
Nicolas Patry's avatar
Nicolas Patry committed
14
use tracing::warn;
15
use utoipa::ToSchema;
Olivier Dehaene's avatar
Olivier Dehaene committed
16
use validation::Validation;
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
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
#[derive(PartialEq)]
pub enum Attention {
    Paged,
    FlashDecoding,
    FlashInfer,
}

#[derive(Debug)]
pub struct ParseError;

impl std::fmt::Display for ParseError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "Cannot parse attention value")
    }
}
impl std::error::Error for ParseError {}

impl std::str::FromStr for Attention {
    type Err = ParseError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "paged" => Ok(Attention::Paged),
            "flashdecoding" => Ok(Attention::FlashDecoding),
            "flashinfer" => Ok(Attention::FlashInfer),
            _ => Err(ParseError),
        }
    }
}

drbh's avatar
drbh committed
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
#[derive(Clone, Deserialize, ToSchema)]
pub(crate) struct VertexInstance {
    #[schema(example = "What is Deep Learning?")]
    pub inputs: String,
    #[schema(nullable = true, default = "null", example = "null")]
    pub parameters: Option<GenerateParameters>,
}

#[derive(Deserialize, ToSchema)]
pub(crate) struct VertexRequest {
    #[serde(rename = "instances")]
    pub instances: Vec<VertexInstance>,
}

#[derive(Clone, Deserialize, ToSchema, Serialize)]
pub(crate) struct VertexResponse {
    pub predictions: Vec<String>,
}

66
67
/// Hub type
#[derive(Clone, Debug, Deserialize)]
68
pub struct HubModelInfo {
69
70
71
72
73
74
    #[serde(rename(deserialize = "id"))]
    pub model_id: String,
    pub sha: Option<String>,
    pub pipeline_tag: Option<String>,
}

75
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
76
77
78
79
80
pub struct ChatTemplate {
    name: String,
    template: String,
}

81
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
82
83
84
85
86
87
#[serde(untagged)]
pub enum ChatTemplateVersions {
    Single(String),
    Multiple(Vec<ChatTemplate>),
}

88
89
use std::path::Path;

90
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
91
pub struct HubTokenizerConfig {
92
    pub chat_template: Option<ChatTemplateVersions>,
93
    pub completion_template: Option<String>,
94
95
    pub bos_token: Option<TokenizerConfigToken>,
    pub eos_token: Option<TokenizerConfigToken>,
96
97
98
    pub tokenizer_class: Option<String>,
    pub add_bos_token: Option<bool>,
    pub add_eos_token: Option<bool>,
99
100
101
}

impl HubTokenizerConfig {
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
    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,
        }
122
123
124
    }
}

125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
#[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
144
145
146
147
148
149
150
151
#[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 {
152
153
154
155
    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
156
157
158
    }
}

159
#[derive(Clone, Debug, Deserialize, ToSchema, Serialize)]
drbh's avatar
drbh committed
160
161
#[serde(tag = "type", content = "value")]
pub(crate) enum GrammarType {
162
163
164
165
166
    /// 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
167
    #[serde(alias = "json_object")]
168
169
    #[schema(example = json ! ({"properties": {"location":{"type": "string"}}}))]
    Json(serde_json::Value),
drbh's avatar
drbh committed
170
171
172
173
    #[serde(rename = "regex")]
    Regex(String),
}

174
175
#[derive(Clone, Debug, Serialize, ToSchema)]
pub struct Info {
176
    /// Model info
177
178
179
180
    #[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
181
182
183
184
    // #[schema(example = "torch.float16")]
    // pub model_dtype: String,
    // #[schema(example = "cuda")]
    // pub model_device_type: String,
185
186
    #[schema(nullable = true, example = "text-generation")]
    pub model_pipeline_tag: Option<String>,
Nicolas Patry's avatar
Nicolas Patry committed
187

188
189
190
191
192
193
194
195
    /// 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
196
    pub max_input_tokens: usize,
197
198
199
200
    #[schema(example = "2048")]
    pub max_total_tokens: usize,
    #[schema(example = "2")]
    pub validation_workers: usize,
201
202
    #[schema(example = "32")]
    pub max_client_batch_size: usize,
Nicolas Patry's avatar
Nicolas Patry committed
203

204
    /// Router Info
205
206
    #[schema(example = "text-generation-router")]
    pub router: &'static str,
207
208
209
210
    #[schema(example = "0.5.0")]
    pub version: &'static str,
    #[schema(nullable = true, example = "null")]
    pub sha: Option<&'static str>,
211
212
    #[schema(nullable = true, example = "null")]
    pub docker_label: Option<&'static str>,
213
214
}

drbh's avatar
drbh committed
215
#[derive(Clone, Debug, Deserialize, ToSchema, Default)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
216
pub(crate) struct GenerateParameters {
217
    /// Generate best_of sequences and return the one if the highest token logprobs.
218
219
220
    #[serde(default)]
    #[schema(exclusive_minimum = 0, nullable = true, default = "null", example = 1)]
    pub best_of: Option<usize>,
221
222

    /// The value used to module the logits distribution.
223
224
225
226
227
228
229
230
    #[serde(default)]
    #[schema(
        exclusive_minimum = 0.0,
        nullable = true,
        default = "null",
        example = 0.5
    )]
    pub temperature: Option<f32>,
231
232
233

    /// The parameter for repetition penalty. 1.0 means no penalty.
    /// See [this paper](https://arxiv.org/pdf/1909.05858.pdf) for more details.
234
235
236
237
238
239
240
241
    #[serde(default)]
    #[schema(
        exclusive_minimum = 0.0,
        nullable = true,
        default = "null",
        example = 1.03
    )]
    pub repetition_penalty: Option<f32>,
242
243
244
245

    /// 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.
246
    #[serde(default)]
247
248
249
250
251
252
253
    #[schema(
        exclusive_minimum = -2.0,
        nullable = true,
        default = "null",
        example = 0.1
    )]
    pub frequency_penalty: Option<f32>,
254
255

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

    /// Top-p value for nucleus sampling.
261
262
263
264
265
266
267
268
269
    #[serde(default)]
    #[schema(
        exclusive_minimum = 0.0,
        maximum = 1.0,
        nullable = true,
        default = "null",
        example = 0.95
    )]
    pub top_p: Option<f32>,
270
271
272

    /// Typical Decoding mass
    /// See [Typical Decoding for Natural Language Generation](https://arxiv.org/abs/2202.00666) for more information.
273
    #[serde(default)]
274
275
276
277
278
279
280
281
    #[schema(
        exclusive_minimum = 0.0,
        maximum = 1.0,
        nullable = true,
        default = "null",
        example = 0.95
    )]
    pub typical_p: Option<f32>,
282
283

    /// Activate logits sampling.
284
    #[serde(default)]
285
    #[schema(default = "false", example = true)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
286
    pub do_sample: bool,
287
288

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

    /// Whether to prepend the prompt to the generated text
OlivierDehaene's avatar
OlivierDehaene committed
294
    #[serde(default)]
295
    #[schema(nullable = true, default = "null", example = false)]
296
    pub return_full_text: Option<bool>,
297
298

    /// Stop generating tokens if a member of `stop` is generated.
299
    #[serde(default)]
300
    #[schema(inline, max_items = 4, example = json ! (["photographer"]))]
301
    pub stop: Vec<String>,
302
303

    /// Truncate inputs tokens to the given size.
OlivierDehaene's avatar
OlivierDehaene committed
304
    #[serde(default)]
305
    #[schema(nullable = true, default = "null", example = "null")]
306
    pub truncate: Option<usize>,
307
308

    /// Watermarking with [A Watermark for Large Language Models](https://arxiv.org/abs/2301.10226).
309
    #[serde(default)]
310
311
    #[schema(default = "false", example = true)]
    pub watermark: bool,
312
313

    /// Whether to return generation details.
314
    #[serde(default)]
315
    #[schema(default = "true")]
OlivierDehaene's avatar
OlivierDehaene committed
316
    pub details: bool,
317
318

    /// Whether to return decoder input token logprobs and ids.
319
    #[serde(default)]
320
    #[schema(default = "false")]
321
    pub decoder_input_details: bool,
322
323

    /// Random sampling seed.
324
    #[serde(default)]
325
326
327
328
329
330
    #[schema(
        exclusive_minimum = 0,
        nullable = true,
        default = "null",
        example = "null"
    )]
331
    pub seed: Option<u64>,
332
333

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

    /// Grammar constraints for the generation.
drbh's avatar
drbh committed
339
    #[serde(default)]
340
    #[schema(nullable = true, default = "null", example = "null")]
drbh's avatar
drbh committed
341
    pub grammar: Option<GrammarType>,
drbh's avatar
drbh committed
342
343
344
345
346

    /// 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
347
348
}

349
fn default_max_new_tokens() -> Option<u32> {
350
    Some(100)
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
351
352
353
354
}

fn default_parameters() -> GenerateParameters {
    GenerateParameters {
355
        best_of: None,
356
357
        temperature: None,
        repetition_penalty: None,
358
        frequency_penalty: None,
359
360
        top_k: None,
        top_p: None,
361
        typical_p: None,
362
        do_sample: true,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
363
        max_new_tokens: default_max_new_tokens(),
364
        return_full_text: None,
365
        stop: Vec::new(),
366
        truncate: None,
367
        watermark: false,
OlivierDehaene's avatar
OlivierDehaene committed
368
        details: false,
369
        decoder_input_details: false,
370
        seed: None,
Nicolas Patry's avatar
Nicolas Patry committed
371
        top_n_tokens: None,
drbh's avatar
drbh committed
372
        grammar: None,
drbh's avatar
drbh committed
373
        adapter_id: None,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
374
375
376
    }
}

377
378
379
380
381
382
383
384
385
386
387
388
389
#[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;
390

391
    fn try_from(value: PromptDeserializer) -> Result<Self, Self::Error> {
392
        match value {
393
394
395
396
397
398
399
400
401
402
403
            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))
                }
            }
404
405
406
407
        }
    }
}

408
409
410
411
412
#[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.
413
    pub model: Option<String>,
414
415
416

    /// The prompt to generate completions for.
    #[schema(example = "What is Deep Learning?")]
417
    pub prompt: Prompt,
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454

    /// 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>,
455
456
457
458
459

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

462
463
464
465
466
467
468
469
470
#[derive(Clone, Serialize, ToSchema)]
#[serde(tag = "object")]
enum Completion {
    #[serde(rename = "text_completion")]
    Chunk(Chunk),
    #[serde(rename = "text_completion")]
    Final(CompletionFinal),
}

471
#[derive(Clone, Deserialize, Serialize, ToSchema, Default)]
472
pub(crate) struct CompletionFinal {
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
    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,
}

491
492
493
494
495
496
497
498
499
#[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,
}

500
#[derive(Clone, Deserialize, Serialize, ToSchema)]
501
502
pub(crate) struct ChatCompletion {
    pub id: String,
503
    #[schema(example = "1706270835")]
504
    pub created: u64,
505
    #[schema(example = "mistralai/Mistral-7B-Instruct-v0.2")]
506
507
508
509
510
511
    pub model: String,
    pub system_fingerprint: String,
    pub choices: Vec<ChatCompletionComplete>,
    pub usage: Usage,
}

512
#[derive(Clone, Deserialize, Serialize, ToSchema)]
513
514
pub(crate) struct ChatCompletionComplete {
    pub index: u32,
Nicolas Patry's avatar
Nicolas Patry committed
515
    pub message: OutputMessage,
516
    pub logprobs: Option<ChatCompletionLogprobs>,
517
518
519
    pub finish_reason: String,
}

520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
#[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;
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562

        // 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
563
564
565
566
567
568
                        .into_iter()
                        .map(|t| ChatCompletionTopLogprob {
                            token: t.text,
                            logprob: t.logprob,
                        })
                        .collect(),
569
570
571
572
573
574
                    None => vec![], // Handle the case where there are no top tokens
                },
            })
            .collect();

        Self { content }
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
    }
}

#[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,
}

591
#[derive(Clone, Deserialize, Serialize, ToSchema, Default)]
592
593
594
595
596
597
pub(crate) struct Usage {
    pub prompt_tokens: u32,
    pub completion_tokens: u32,
    pub total_tokens: u32,
}

598
599
600
601
602
603
604
605
606
#[derive(Clone, Serialize, ToSchema)]
#[serde(tag = "object")]
enum CompletionType {
    #[serde(rename = "chat.completion.chunk")]
    ChatCompletionChunk(ChatCompletionChunk),
    #[serde(rename = "chat.completion")]
    ChatCompletion(ChatCompletion),
}

607
608
609
610
impl ChatCompletion {
    pub(crate) fn new(
        model: String,
        system_fingerprint: String,
drbh's avatar
drbh committed
611
        output: Option<String>,
612
613
614
        created: u64,
        details: Details,
        return_logprobs: bool,
615
        tool_calls: Option<Vec<ToolCall>>,
616
    ) -> Self {
Nicolas Patry's avatar
Nicolas Patry committed
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
        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(),
                })
            }
        };
641
642
643
644
645
646
647
        Self {
            id: String::new(),
            created,
            model,
            system_fingerprint,
            choices: vec![ChatCompletionComplete {
                index: 0,
Nicolas Patry's avatar
Nicolas Patry committed
648
                message,
649
                logprobs: return_logprobs
650
                    .then(|| ChatCompletionLogprobs::from((details.tokens, details.top_tokens))),
651
                finish_reason: details.finish_reason.format(true),
652
653
654
655
656
657
658
659
660
            }],
            usage: Usage {
                prompt_tokens: details.prefill.len() as u32,
                completion_tokens: details.generated_tokens,
                total_tokens: details.prefill.len() as u32 + details.generated_tokens,
            },
        }
    }
}
661
#[derive(Clone, Serialize, ToSchema)]
662
663
pub(crate) struct ChatCompletionChunk {
    pub id: String,
664
    #[schema(example = "1706270978")]
665
    pub created: u64,
666
    #[schema(example = "mistralai/Mistral-7B-Instruct-v0.2")]
667
668
669
670
671
    pub model: String,
    pub system_fingerprint: String,
    pub choices: Vec<ChatCompletionChoice>,
}

672
#[derive(Clone, Serialize, ToSchema)]
673
674
675
pub(crate) struct ChatCompletionChoice {
    pub index: u32,
    pub delta: ChatCompletionDelta,
676
    pub logprobs: Option<ChatCompletionLogprobs>,
677
678
679
    pub finish_reason: Option<String>,
}

Nicolas Patry's avatar
Nicolas Patry committed
680
681
682
683
684
685
686
#[derive(Clone, Deserialize, ToSchema, Serialize, Debug, PartialEq)]
pub struct ToolCallDelta {
    #[schema(example = "assistant")]
    role: String,
    tool_calls: DeltaToolCall,
}

687
688
#[derive(Clone, Debug, Serialize, ToSchema)]
#[serde(untagged)]
Nicolas Patry's avatar
Nicolas Patry committed
689
690
691
enum ChatCompletionDelta {
    Chat(TextMessage),
    Tool(ToolCallDelta),
drbh's avatar
drbh committed
692
693
}

Nicolas Patry's avatar
Nicolas Patry committed
694
#[derive(Clone, Deserialize, Serialize, ToSchema, Debug, PartialEq)]
drbh's avatar
drbh committed
695
696
697
698
699
700
701
pub(crate) struct DeltaToolCall {
    pub index: u32,
    pub id: String,
    pub r#type: String,
    pub function: Function,
}

Nicolas Patry's avatar
Nicolas Patry committed
702
#[derive(Clone, Deserialize, Serialize, ToSchema, Debug, PartialEq)]
drbh's avatar
drbh committed
703
704
705
pub(crate) struct Function {
    pub name: Option<String>,
    pub arguments: String,
706
707
}

drbh's avatar
drbh committed
708
#[allow(clippy::too_many_arguments)]
709
710
711
712
impl ChatCompletionChunk {
    pub(crate) fn new(
        model: String,
        system_fingerprint: String,
drbh's avatar
drbh committed
713
714
        delta: Option<String>,
        tool_calls: Option<Vec<String>>,
715
        created: u64,
716
        logprobs: Option<ChatCompletionLogprobs>,
717
718
        finish_reason: Option<String>,
    ) -> Self {
719
        let delta = match (delta, tool_calls) {
Nicolas Patry's avatar
Nicolas Patry committed
720
721
722
723
724
725
726
            (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 {
727
728
729
730
731
732
733
                    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
734
735
736
737
738
739
                },
            }),
            (None, None) => ChatCompletionDelta::Chat(TextMessage {
                role: "assistant".to_string(),
                content: "".to_string(),
            }),
740
        };
741
742
743
744
745
746
        Self {
            id: String::new(),
            created,
            model,
            system_fingerprint,
            choices: vec![ChatCompletionChoice {
747
                index: 0,
748
                delta,
749
750
751
752
753
754
755
756
757
                logprobs,
                finish_reason,
            }],
        }
    }
}

#[derive(Clone, Deserialize, ToSchema, Serialize)]
pub(crate) struct ChatRequest {
758
    #[schema(example = "mistralai/Mistral-7B-Instruct-v0.2")]
drbh's avatar
drbh committed
759
    /// [UNUSED] ID of the model to use. See the model endpoint compatibility table for details on which models work with the Chat API.
760
    pub model: Option<String>,
drbh's avatar
drbh committed
761

762
    /// A list of messages comprising the conversation so far.
drbh's avatar
drbh committed
763
    #[schema(example = "[{\"role\": \"user\", \"content\": \"What is Deep Learning?\"}]")]
764
765
766
767
768
    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)]
769
    #[schema(example = "1.0")]
770
771
772
773
774
775
776
777
778
779
780
781
782
783
    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)]
784
    #[schema(example = "false")]
785
786
787
788
789
    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)]
790
    #[schema(example = "5")]
791
792
793
794
    pub top_logprobs: Option<u32>,

    /// The maximum number of tokens that can be generated in the chat completion.
    #[serde(default)]
795
    #[schema(example = "32")]
796
797
798
799
800
801
    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)]
802
    #[schema(nullable = true, example = "2")]
803
804
805
806
807
    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)]
808
    #[schema(nullable = true, example = 0.1)]
809
810
    pub presence_penalty: Option<f32>,

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

816
817
818
819
820
    #[serde(default = "bool::default")]
    pub stream: bool,

    #[schema(nullable = true, example = 42)]
    pub seed: Option<u64>,
821
822
823
824
825
826

    /// 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)]
827
    #[schema(nullable = true, example = 1.0)]
828
829
830
831
832
    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)]
833
    #[schema(nullable = true, example = 0.95)]
834
    pub top_p: Option<f32>,
drbh's avatar
drbh committed
835
836
837
838
839
840
841
842

    /// 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
843
    #[serde(default)]
drbh's avatar
drbh committed
844
845
    #[schema(
        nullable = true,
drbh's avatar
drbh committed
846
        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
847
848
849
850
851
852
    )]
    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
853
    pub tool_choice: ToolChoice,
drbh's avatar
drbh committed
854
855
856
857
858
859
860

    /// 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>,
861
862
863
864
865

    /// A guideline to be used in the chat_template
    #[serde(default)]
    #[schema(nullable = true, default = "null", example = "null")]
    pub guideline: Option<String>,
drbh's avatar
drbh committed
866
867
}

drbh's avatar
drbh committed
868
869
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
870
}
871
872
873
874

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize, ToSchema)]
#[serde(untagged)]
pub enum ToolType {
drbh's avatar
drbh committed
875
    OneOf,
876
877
    FunctionName(String),
    Function { function: FunctionName },
drbh's avatar
drbh committed
878
    NoTool,
drbh's avatar
drbh committed
879
880
}

881
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema)]
882
883
884
885
pub struct FunctionName {
    pub name: String,
}

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

890
891
892
#[derive(Deserialize)]
#[serde(untagged)]
enum ToolTypeDeserializer {
drbh's avatar
drbh committed
893
894
    String(String),
    ToolType(ToolType),
895
}
drbh's avatar
drbh committed
896

897
898
impl From<ToolTypeDeserializer> for ToolChoice {
    fn from(value: ToolTypeDeserializer) -> Self {
drbh's avatar
drbh committed
899
        match value {
drbh's avatar
drbh committed
900
901
902
903
            ToolTypeDeserializer::String(s) => match s.as_str() {
                "none" => ToolChoice(Some(ToolType::NoTool)),
                "auto" => ToolChoice(Some(ToolType::OneOf)),
                _ => ToolChoice(Some(ToolType::FunctionName(s))),
drbh's avatar
drbh committed
904
            },
drbh's avatar
drbh committed
905
            ToolTypeDeserializer::ToolType(tool_type) => ToolChoice(Some(tool_type)),
drbh's avatar
drbh committed
906
907
908
909
        }
    }
}

910
#[derive(Debug, Deserialize, Serialize, ToSchema, PartialEq)]
drbh's avatar
drbh committed
911
pub struct JsonSchemaTool {
drbh's avatar
drbh committed
912
913
914
915
916
    #[serde(flatten)]
    functions_map: FunctionsMap,
    properties: Properties,
}

917
#[derive(Debug, Serialize, Deserialize, PartialEq)]
drbh's avatar
drbh committed
918
919
920
921
922
struct FunctionsMap {
    #[serde(rename = "$functions")]
    functions: std::collections::HashMap<String, serde_json::Value>,
}

923
#[derive(Debug, Serialize, Deserialize, PartialEq)]
drbh's avatar
drbh committed
924
925
926
927
928
struct FunctionRef {
    #[serde(rename = "$ref")]
    ref_path: String,
}

929
#[derive(Debug, Serialize, Deserialize, PartialEq)]
drbh's avatar
drbh committed
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
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
945
#[derive(Clone, Debug, Deserialize, Serialize, ToSchema, Default, PartialEq)]
drbh's avatar
drbh committed
946
947
948
949
pub(crate) struct FunctionDefinition {
    #[serde(default)]
    pub description: Option<String>,
    pub name: String,
950
951
    #[serde(alias = "parameters")]
    pub arguments: serde_json::Value,
drbh's avatar
drbh committed
952
953
954
955
956
957
958
959
960
}

#[derive(Clone, Debug, Deserialize, Serialize, ToSchema)]
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,
961
962
}

963
#[derive(Clone, Serialize, Deserialize, Default)]
964
pub(crate) struct ChatTemplateInputs<'a> {
Nicolas Patry's avatar
Nicolas Patry committed
965
    messages: Vec<TextMessage>,
966
967
    bos_token: Option<&'a str>,
    eos_token: Option<&'a str>,
968
    add_generation_prompt: bool,
drbh's avatar
drbh committed
969
    tools: Option<Vec<Tool>>,
970
    guideline: Option<&'a str>,
971
972
}

Nicolas Patry's avatar
Nicolas Patry committed
973
#[derive(Clone, Deserialize, Serialize, ToSchema, Default, Debug, PartialEq)]
drbh's avatar
drbh committed
974
pub(crate) struct ToolCall {
975
    pub id: String,
drbh's avatar
drbh committed
976
977
978
979
    pub r#type: String,
    pub function: FunctionDefinition,
}

Nicolas Patry's avatar
Nicolas Patry committed
980
#[derive(Clone, Deserialize, ToSchema, Serialize, Debug, PartialEq)]
981
pub struct Url {
Nicolas Patry's avatar
Nicolas Patry committed
982
    url: String,
drbh's avatar
drbh committed
983
984
}

Nicolas Patry's avatar
Nicolas Patry committed
985
986
987
#[derive(Clone, Deserialize, ToSchema, Serialize, Debug, PartialEq)]
#[serde(tag = "type")]
#[serde(rename_all = "snake_case")]
988
989
990
pub enum MessageChunk {
    Text { text: String },
    ImageUrl { image_url: Url },
Nicolas Patry's avatar
Nicolas Patry committed
991
992
993
994
995
996
997
}

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

1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
#[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
1016
1017
1018
1019
                *self = MessageContent::MultipleChunks(vec![
                    MessageChunk::Text { text: text.clone() },
                    chunk,
                ]);
Nicolas Patry's avatar
Nicolas Patry committed
1020
            }
1021
1022
1023
1024
            MessageContent::MultipleChunks(chunks) => {
                chunks.push(chunk);
            }
        }
drbh's avatar
drbh committed
1025
1026
1027
    }
}

Nicolas Patry's avatar
Nicolas Patry committed
1028
1029
#[derive(Clone, Deserialize, ToSchema, Serialize, Debug, PartialEq)]
pub struct TextMessage {
1030
1031
1032
    #[schema(example = "user")]
    pub role: String,
    #[schema(example = "My name is David and I")]
Nicolas Patry's avatar
Nicolas Patry committed
1033
1034
1035
1036
1037
1038
1039
    pub content: String,
}

impl From<Message> for TextMessage {
    fn from(value: Message) -> Self {
        TextMessage {
            role: value.role,
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
            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
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
        }
    }
}

#[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),
1067
1068
}

1069
#[derive(Clone, Debug, Deserialize, ToSchema)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1070
pub(crate) struct GenerateRequest {
1071
    #[schema(example = "My name is Olivier and I")]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1072
1073
1074
1075
1076
    pub inputs: String,
    #[serde(default = "default_parameters")]
    pub parameters: GenerateParameters,
}

1077
1078
1079
1080
1081
1082
1083
#[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
1084
    #[schema(default = "false")]
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
    pub stream: bool,
}

impl From<CompatGenerateRequest> for GenerateRequest {
    fn from(req: CompatGenerateRequest) -> Self {
        Self {
            inputs: req.inputs,
            parameters: req.parameters,
        }
    }
}

1097
1098
1099
#[derive(Debug, Serialize, ToSchema)]
pub struct PrefillToken {
    #[schema(example = 0)]
Nicolas Patry's avatar
Nicolas Patry committed
1100
    pub id: u32,
1101
    #[schema(example = "test")]
Nicolas Patry's avatar
Nicolas Patry committed
1102
    pub text: String,
1103
    #[schema(nullable = true, example = - 0.34)]
Nicolas Patry's avatar
Nicolas Patry committed
1104
    pub logprob: f32,
1105
1106
}

1107
#[derive(Debug, Serialize, ToSchema, Clone)]
1108
1109
pub struct Token {
    #[schema(example = 0)]
Nicolas Patry's avatar
Nicolas Patry committed
1110
    pub id: u32,
1111
    #[schema(example = "test")]
Nicolas Patry's avatar
Nicolas Patry committed
1112
    pub text: String,
1113
    #[schema(nullable = true, example = - 0.34)]
Nicolas Patry's avatar
Nicolas Patry committed
1114
    pub logprob: f32,
1115
    #[schema(example = "false")]
Nicolas Patry's avatar
Nicolas Patry committed
1116
    pub special: bool,
1117
1118
}

1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
#[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
1131
#[derive(Debug, Serialize, ToSchema)]
1132
#[serde(rename_all(serialize = "snake_case"))]
1133
#[schema(example = "Length")]
Nicolas Patry's avatar
Nicolas Patry committed
1134
pub enum FinishReason {
1135
1136
1137
1138
1139
1140
1141
1142
    #[schema(rename = "length")]
    Length,
    #[serde(rename = "eos_token")]
    #[schema(rename = "eos_token")]
    EndOfSequenceToken,
    #[schema(rename = "stop_sequence")]
    StopSequence,
}
1143

1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
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"),
        }
    }
}

1154
1155
1156
1157
1158
1159
1160
1161
1162
impl FinishReason {
    pub fn format(&self, use_stop: bool) -> String {
        match self {
            FinishReason::EndOfSequenceToken if use_stop => "stop".to_string(),
            _ => self.to_string(),
        }
    }
}

1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
#[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
1175
1176
    #[serde(skip_serializing_if = "Vec::is_empty")]
    pub top_tokens: Vec<Vec<Token>>,
1177
1178
}

1179
#[derive(Serialize, ToSchema)]
OlivierDehaene's avatar
OlivierDehaene committed
1180
pub(crate) struct Details {
1181
1182
1183
    #[schema(example = "length")]
    pub finish_reason: FinishReason,
    #[schema(example = 1)]
OlivierDehaene's avatar
OlivierDehaene committed
1184
    pub generated_tokens: u32,
1185
    #[schema(nullable = true, example = 42)]
1186
    pub seed: Option<u64>,
1187
1188
    pub prefill: Vec<PrefillToken>,
    pub tokens: Vec<Token>,
1189
1190
    #[serde(skip_serializing_if = "Option::is_none")]
    pub best_of_sequences: Option<Vec<BestOfSequence>>,
Nicolas Patry's avatar
Nicolas Patry committed
1191
1192
    #[serde(skip_serializing_if = "Vec::is_empty")]
    pub top_tokens: Vec<Vec<Token>>,
OlivierDehaene's avatar
OlivierDehaene committed
1193
1194
}

1195
#[derive(Serialize, ToSchema)]
1196
pub(crate) struct GenerateResponse {
1197
    #[schema(example = "test")]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1198
    pub generated_text: String,
OlivierDehaene's avatar
OlivierDehaene committed
1199
1200
    #[serde(skip_serializing_if = "Option::is_none")]
    pub details: Option<Details>,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1201
}
1202

1203
1204
1205
1206
1207
1208
#[derive(Serialize, ToSchema)]
pub(crate) struct ChatTokenizeResponse {
    pub(crate) tokenize_response: TokenizeResponse,
    pub(crate) templated_text: String,
}

1209
1210
1211
1212
#[derive(Serialize, ToSchema)]
#[serde(transparent)]
pub(crate) struct TokenizeResponse(Vec<SimpleToken>);

1213
1214
1215
1216
1217
1218
#[derive(Serialize, ToSchema)]
pub(crate) struct StreamDetails {
    #[schema(example = "length")]
    pub finish_reason: FinishReason,
    #[schema(example = 1)]
    pub generated_tokens: u32,
1219
    #[schema(nullable = true, example = 42)]
1220
    pub seed: Option<u64>,
1221
1222
    #[schema(example = 1)]
    pub input_length: u32,
1223
1224
1225
}

#[derive(Serialize, ToSchema)]
1226
pub(crate) struct StreamResponse {
1227
    pub index: u32,
1228
    pub token: Token,
Nicolas Patry's avatar
Nicolas Patry committed
1229
1230
    #[serde(skip_serializing_if = "Vec::is_empty")]
    pub top_tokens: Vec<Token>,
1231
    #[schema(nullable = true, default = "null", example = "test")]
1232
    pub generated_text: Option<String>,
1233
1234
    #[schema(nullable = true, default = "null")]
    pub details: Option<StreamDetails>,
1235
1236
}

1237
#[derive(Serialize, ToSchema)]
1238
1239
pub(crate) struct ErrorResponse {
    pub error: String,
1240
    pub error_type: String,
1241
}
1242
1243

#[cfg(test)]
1244
mod tests {
1245
    use super::*;
Nicolas Patry's avatar
Nicolas Patry committed
1246
    use serde_json::json;
1247
1248
    use tokenizers::Tokenizer;

1249
    pub(crate) async fn get_tokenizer() -> Tokenizer {
1250
1251
1252
1253
        let api = hf_hub::api::sync::Api::new().unwrap();
        let repo = api.model("gpt2".to_string());
        let filename = repo.get("tokenizer.json").unwrap();
        Tokenizer::from_file(filename).unwrap()
1254
    }
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268

    #[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
1269
1270
1271
1272
        assert_eq!(
            config.chat_template,
            Some(ChatTemplateVersions::Single("test".to_string()))
        );
1273
1274
        assert_eq!(
            config.bos_token,
1275
1276
1277
1278
1279
1280
1281
1282
1283
            Some(TokenizerConfigToken::String(
                "<|begin▁of▁sentence|>".to_string()
            ))
        );
        assert_eq!(
            config.eos_token,
            Some(TokenizerConfigToken::String(
                "<|end▁of▁sentence|>".to_string()
            ))
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
        );

        // 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
1311
1312
1313
1314
        assert_eq!(
            config.chat_template,
            Some(ChatTemplateVersions::Single("test".to_string()))
        );
1315
1316
        assert_eq!(
            config.bos_token,
1317
1318
1319
1320
1321
1322
1323
1324
1325
            Some(TokenizerConfigToken::Object {
                content: "<|begin▁of▁sentence|>".to_string()
            })
        );
        assert_eq!(
            config.eos_token,
            Some(TokenizerConfigToken::Object {
                content: "<|end▁of▁sentence|>".to_string()
            })
1326
1327
        );
    }
Nicolas Patry's avatar
Nicolas Patry committed
1328
1329
1330

    #[test]
    fn test_chat_simple_string() {
Nicolas Patry's avatar
Nicolas Patry committed
1331
        let json = json!({
Nicolas Patry's avatar
Nicolas Patry committed
1332
            "model": "",
Nicolas Patry's avatar
Nicolas Patry committed
1333
1334
            "messages": [{
                "role": "user",
Nicolas Patry's avatar
Nicolas Patry committed
1335
                "content": "What is Deep Learning?"
Nicolas Patry's avatar
Nicolas Patry committed
1336
            }]
Nicolas Patry's avatar
Nicolas Patry committed
1337
1338
1339
1340
1341
1342
1343
        });
        let request: ChatRequest = serde_json::from_str(json.to_string().as_str()).unwrap();

        assert_eq!(
            request.messages[0],
            Message {
                role: "user".to_string(),
1344
                content: MessageContent::SingleText("What is Deep Learning?".to_string()),
Nicolas Patry's avatar
Nicolas Patry committed
1345
1346
1347
1348
1349
                name: None
            }
        );
    }

drbh's avatar
drbh committed
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
    #[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
1379
1380
    #[test]
    fn test_chat_request() {
Nicolas Patry's avatar
Nicolas Patry committed
1381
        let json = json!({
Nicolas Patry's avatar
Nicolas Patry committed
1382
            "model": "",
Nicolas Patry's avatar
Nicolas Patry committed
1383
1384
            "messages": [{
                "role": "user",
Nicolas Patry's avatar
Nicolas Patry committed
1385
1386
                "content": [
                    {"type": "text", "text": "Whats in this image?"},
Nicolas Patry's avatar
Nicolas Patry committed
1387
                    {"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
1388
                ]
Nicolas Patry's avatar
Nicolas Patry committed
1389
            }]
Nicolas Patry's avatar
Nicolas Patry committed
1390
1391
1392
1393
1394
1395
1396
        });
        let request: ChatRequest = serde_json::from_str(json.to_string().as_str()).unwrap();

        assert_eq!(
            request.messages[0],
            Message{
                role: "user".to_string(),
1397
1398
1399
1400
                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
1401
1402
1403
1404
                name: None
            }
        );
    }
Nicolas Patry's avatar
Nicolas Patry committed
1405
1406
1407
1408
1409

    #[test]
    fn text_message_convert() {
        let message = Message{
                role: "user".to_string(),
1410
1411
1412
1413
                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
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
                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)");
    }
    #[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"}}}]}"#
        );
    }
1451
}