infer.rs 75.9 KB
Newer Older
1
2
/// Batching and inference logic
use crate::validation::{Validation, ValidationError};
3
use crate::{
4
    ChatTemplateInputs, ChatTemplateVersions, Entry, GenerateRequest, GenerateStreamResponse,
Nicolas Patry's avatar
Nicolas Patry committed
5
    HubTokenizerConfig, Message, MessageChunk, PrefillToken, Queue, Text, TextMessage, Token,
6
};
7
use crate::{FunctionRef, FunctionsMap, GrammarType, Properties, Tool, ToolType, Tools};
8
use futures::future::try_join_all;
9
use minijinja::{Environment, ErrorKind, Template};
10
use nohash_hasher::IntMap;
11
12
use serde_json::{json, Map, Value};
use std::collections::HashMap;
13
14
15
16
use std::sync::{
    atomic::{AtomicBool, Ordering},
    Arc,
};
17
use text_generation_client::{
Nicolas Patry's avatar
Nicolas Patry committed
18
    Batch, CachedBatch, ClientError, GeneratedText, Generation, ShardedClient, Tokens,
19
20
};
use thiserror::Error;
OlivierDehaene's avatar
OlivierDehaene committed
21
use tokio::sync::mpsc::error::SendError;
22
use tokio::sync::{mpsc, Notify, Semaphore, TryAcquireError};
23
use tokio::time::Instant;
OlivierDehaene's avatar
OlivierDehaene committed
24
25
use tokio_stream::wrappers::UnboundedReceiverStream;
use tokio_stream::StreamExt;
26
use tracing::{info_span, instrument, Instrument, Span};
27
28
29
30
31
32

/// Inference struct
#[derive(Clone)]
pub struct Infer {
    /// Validation
    validation: Validation,
33
34
    /// Request queue
    queue: Queue,
35
36
    /// Shared state
    shared: Arc<Shared>,
37
38
    /// Chat template
    chat_template: Option<ChatTemplate>,
39
40
41
42
43
44
45
46
47
48
    /// Inference limit
    limit_concurrent_requests: Arc<Semaphore>,
}

/// Infer shared state
struct Shared {
    /// Batching background Tokio task notifier
    batching_task: Notify,
}

49
50
51
52
53
/// Raise a exception (custom function) used in the chat templates
fn raise_exception(err_text: String) -> Result<String, minijinja::Error> {
    Err(minijinja::Error::new(ErrorKind::SyntaxError, err_text))
}

54
impl Infer {
55
    #[allow(clippy::too_many_arguments)]
56
57
58
    pub(crate) fn new(
        client: ShardedClient,
        validation: Validation,
59
        waiting_served_ratio: f32,
60
        max_batch_prefill_tokens: u32,
61
        max_batch_total_tokens: u32,
62
        max_waiting_tokens: usize,
63
        max_batch_size: Option<usize>,
64
        max_concurrent_requests: usize,
65
        requires_padding: bool,
66
        window_size: Option<u32>,
Nicolas Patry's avatar
Nicolas Patry committed
67
        speculate: u32,
68
        generation_health: Arc<AtomicBool>,
69
        tokenizer_config: HubTokenizerConfig,
70
71
    ) -> Self {
        // Infer shared state
Nicolas Patry's avatar
Nicolas Patry committed
72
        let queue = Queue::new(requires_padding, 16, window_size, speculate);
73
74
75
76
77
78
79
        let shared = Arc::new(Shared {
            batching_task: Notify::new(),
        });

        // Spawn batching background task that contains all the inference logic
        tokio::spawn(batching_task(
            client,
80
            waiting_served_ratio,
81
            max_batch_prefill_tokens,
82
            max_batch_total_tokens,
83
            max_waiting_tokens,
84
            max_batch_size,
85
            queue.clone(),
86
            shared.clone(),
87
            generation_health,
88
89
        ));

90
91
        let chat_template = tokenizer_config
            .chat_template
92
93
94
95
96
97
98
99
100
101
102
103
            .and_then(|t| match t {
                ChatTemplateVersions::Single(template) => Some(template),
                ChatTemplateVersions::Multiple(templates) => templates
                    .into_iter()
                    .find(|t| t.name == "default")
                    .map(|t| t.template),
            })
            .map(|t| {
                // .strip() is not supported in minijinja
                let t = t.replace(".strip()", " | trim");
                ChatTemplate::new(t, tokenizer_config.bos_token, tokenizer_config.eos_token)
            });
104

105
106
107
108
109
        // Inference limit with a semaphore
        let semaphore = Arc::new(Semaphore::new(max_concurrent_requests));

        Self {
            validation,
110
            queue,
111
            shared,
112
            chat_template,
113
114
115
116
            limit_concurrent_requests: semaphore,
        }
    }

117
    /// Add a new request to the queue and return a stream of InferStreamResponse
118
    #[instrument(skip_all)]
119
120
121
    pub(crate) async fn generate_stream(
        &self,
        request: GenerateRequest,
122
    ) -> Result<GenerateStreamResponse, InferError> {
123
        // Limit concurrent requests by acquiring a permit from the semaphore
124
125
126
127
128
        let permit = self
            .clone()
            .limit_concurrent_requests
            .try_acquire_owned()
            .map_err(|err| {
129
                metrics::increment_counter!("tgi_request_failure", "err" => "overloaded");
130
131
132
                tracing::error!("{err}");
                err
            })?;
133
134

        // Validate request
135
136
137
138
139
        let valid_request = self.validation.validate(request).await.map_err(|err| {
            metrics::increment_counter!("tgi_request_failure", "err" => "validation");
            tracing::error!("{err}");
            err
        })?;
140
141

        // MPSC channel to communicate with the background batching task
OlivierDehaene's avatar
OlivierDehaene committed
142
        let (response_tx, response_rx) = mpsc::unbounded_channel();
143
        let input_length = valid_request.input_length;
144

145
146
        // Append the request to the queue
        self.queue.append(Entry {
147
148
            request: valid_request,
            response_tx,
149
150
151
            span: Span::current(),
            temp_span: None,
            queue_time: Instant::now(),
152
153
154
            batch_time: None,
        });

155
        // Notify the background task that we have a new entry in the queue that needs
156
157
158
159
        // to be batched
        self.shared.batching_task.notify_one();

        // Return stream
160
161
162
163
164
        Ok((
            permit,
            input_length,
            UnboundedReceiverStream::new(response_rx),
        ))
165
166
    }

167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
    /// Tokenizer the input
    #[instrument(skip_all)]
    pub(crate) async fn tokenize(
        &self,
        request: GenerateRequest,
    ) -> Result<Option<tokenizers::Encoding>, InferError> {
        // Tokenize request
        let inputs = request.inputs;
        let truncate = request.parameters.truncate;
        let encoding = self
            .validation
            .tokenize(inputs, truncate)
            .await
            .map_err(|err| {
                tracing::error!("Tokenization {err}");
                err
            })?;

        // Return Encoding
        Ok(encoding.map(|(encoding, _)| encoding))
    }

189
190
    /// Apply the chat template to the chat request
    #[instrument(skip_all)]
191
192
193
194
195
    pub(crate) fn apply_chat_template(
        &self,
        messages: Vec<Message>,
        grammar_with_prompt: Option<(GrammarType, String)>,
    ) -> Result<String, InferError> {
196
        self.chat_template
197
198
            .as_ref()
            .ok_or_else(|| InferError::TemplateError(ErrorKind::TemplateNotFound.into()))?
199
            .apply(messages, grammar_with_prompt)
200
201
202
            .map_err(|e| {
                metrics::increment_counter!("tgi_request_failure", "err" => "template");
                tracing::error!("{e}");
203
                e
204
205
206
            })
    }

207
    /// Add a new request to the queue and return a InferResponse
208
    #[instrument(skip_all)]
209
210
211
212
    pub(crate) async fn generate(
        &self,
        request: GenerateRequest,
    ) -> Result<InferResponse, InferError> {
Nicolas Patry's avatar
Nicolas Patry committed
213
214
        let use_top_tokens = request.parameters.top_n_tokens.is_some_and(|x| x > 0);

215
        // Create stream and keep semaphore permit as long as generate lives
216
        let (_permit, _input_length, mut stream) = self.generate_stream(request).await?;
217
218
219
220

        // Return values
        let mut result_prefill = Vec::new();
        let mut result_tokens = Vec::new();
Nicolas Patry's avatar
Nicolas Patry committed
221
        let mut result_top_tokens = Vec::new();
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
        let mut result_generated_text = None;
        let mut result_start = None;
        let mut result_queued = None;

        // Iterate on stream
        while let Some(response) = stream.next().await {
            match response? {
                // Add prefill tokens
                InferStreamResponse::Prefill(tokens) => {
                    // Create Token objects
                    // We do that here instead of in the Python code as Rust for loops are faster
                    result_prefill = tokens
                        .ids
                        .into_iter()
                        .zip(tokens.logprobs.into_iter())
                        .zip(tokens.texts.into_iter())
238
                        .map(|((id, logprob), text)| PrefillToken { id, text, logprob })
239
240
241
                        .collect();
                }
                // Push last token
Nicolas Patry's avatar
Nicolas Patry committed
242
243
244
245
                InferStreamResponse::Intermediate { token, top_tokens } => {
                    result_tokens.push(token);
                    result_top_tokens.push(top_tokens);
                }
246
247
248
249
250
251
252
                // Final message
                // Set return values
                InferStreamResponse::End {
                    token,
                    generated_text,
                    start,
                    queued,
Nicolas Patry's avatar
Nicolas Patry committed
253
                    top_tokens,
254
255
                } => {
                    result_tokens.push(token);
Nicolas Patry's avatar
Nicolas Patry committed
256
                    result_top_tokens.push(top_tokens);
257
258
259
260
261
262
263
264
265
266
267
268
269
                    result_generated_text = Some(generated_text);
                    result_start = Some(start);
                    result_queued = Some(queued)
                }
            }
        }

        // Check that we received a `InferStreamResponse::End` message
        if let (Some(generated_text), Some(queued), Some(start)) =
            (result_generated_text, result_queued, result_start)
        {
            Ok(InferResponse {
                prefill: result_prefill,
270
                _input_length,
271
272
273
274
                tokens: result_tokens,
                generated_text,
                queued,
                start,
Nicolas Patry's avatar
Nicolas Patry committed
275
276
277
278
279
                top_tokens: if use_top_tokens {
                    result_top_tokens
                } else {
                    Vec::new()
                },
280
281
            })
        } else {
282
            let err = InferError::IncompleteGeneration;
283
            metrics::increment_counter!("tgi_request_failure", "err" => "incomplete");
284
285
            tracing::error!("{err}");
            Err(err)
286
287
        }
    }
288
289
    /// Add best_of new requests to the queue and return a InferResponse of the sequence with
    /// the highest log probability per token
290
    #[instrument(skip(self, request))]
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
    pub(crate) async fn generate_best_of(
        &self,
        request: GenerateRequest,
        best_of: usize,
    ) -> Result<(InferResponse, Vec<InferResponse>), InferError> {
        // validate  best_of parameter separately
        let best_of = self.validation.validate_best_of(best_of)?;

        // create multiple generate requests
        let mut infer_responses: Vec<InferResponse> =
            try_join_all((0..best_of).map(|_| self.generate(request.clone()))).await?;

        // get the sequence with the highest log probability per token
        let mut max_index = 0;
        let mut max_logprob: f32 = f32::MIN;

        for (i, response) in infer_responses.iter().enumerate() {
            // mean logprobs of the generated tokens
            let sequence_logprob = response
                .tokens
                .iter()
                .map(|token| token.logprob)
                .sum::<f32>()
                / response.tokens.len() as f32;

            // set best sequence
            if sequence_logprob > max_logprob {
                max_index = i;
                max_logprob = sequence_logprob;
            }
        }
        let best_response = infer_responses.remove(max_index);
        Ok((best_response, infer_responses))
    }
325
326
}

327
328
329
330
331
#[derive(Clone)]
struct ChatTemplate {
    template: Template<'static, 'static>,
    bos_token: Option<String>,
    eos_token: Option<String>,
332
    use_default_tool_template: bool,
333
334
335
336
337
338
339
}

impl ChatTemplate {
    fn new(template: String, bos_token: Option<String>, eos_token: Option<String>) -> Self {
        let mut env = Box::new(Environment::new());
        let template_str = template.into_boxed_str();
        env.add_function("raise_exception", raise_exception);
340
341
342
343

        // check if contains the tools variable within the template
        let use_default_tool_template =
            !template_str.as_ref().replace(' ', "").contains("{{tools}}");
344
345
346
347
348
349
350
351
352
        // leaking env and template_str as read-only, static resources for performance.
        let template = Box::leak(env)
            .template_from_str(Box::leak(template_str))
            .unwrap();

        Self {
            template,
            bos_token,
            eos_token,
353
            use_default_tool_template,
354
355
356
        }
    }

357
358
359
360
361
362
363
364
    fn apply(
        &self,
        mut messages: Vec<Message>,
        grammar_with_prompt: Option<(GrammarType, String)>,
    ) -> Result<String, InferError> {
        if self.use_default_tool_template {
            if let Some(last_message) = messages.last_mut() {
                if let Some((GrammarType::Json(tools), tool_prompt)) = grammar_with_prompt {
Nicolas Patry's avatar
Nicolas Patry committed
365
366
367
                    last_message.content.push(MessageChunk::Text(Text {
                        text: format!("\n---\n{}\n{}", tool_prompt, tools),
                    }));
368
369
370
371
                }
            }
        }

Nicolas Patry's avatar
Nicolas Patry committed
372
373
        let messages: Vec<TextMessage> = messages.into_iter().map(|c| c.into()).collect();

374
375
376
377
378
379
        self.template
            .render(ChatTemplateInputs {
                messages,
                bos_token: self.bos_token.as_deref(),
                eos_token: self.eos_token.as_deref(),
                add_generation_prompt: true,
380
381
                tools: None,
                tools_prompt: None,
382
383
384
385
386
            })
            .map_err(InferError::TemplateError)
    }
}

387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
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
455
456
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
494
495
496
497
498
499
500
501
502
503
504
pub struct ToolGrammar {}

impl ToolGrammar {
    pub fn apply(
        tools: Option<Vec<Tool>>,
        tool_choice: Option<ToolType>,
    ) -> Result<Option<Tools>, InferError> {
        if let Some((req_tools, tool_choice)) = tools.zip(tool_choice) {
            // let tool_prompt = tool_prompt.unwrap_or_default();
            let tools_to_use = match tool_choice {
                ToolType::FunctionName(name) => {
                    vec![req_tools
                        .iter()
                        .find(|tool| tool.function.name == *name)
                        .unwrap_or_else(|| panic!("Tool with name {} not found", name))
                        .clone()]
                }
                ToolType::OneOf => req_tools.to_owned(),
            };

            // adds the error notification function for LLM feedback if required
            let mut text_response_properties = Map::new();
            text_response_properties.insert(
                "error".to_string(),
                serde_json::json!({
                    "type": "string",
                    "description": "The error or issue to notify"
                }),
            );
            text_response_properties.insert(
                "_name".to_string(),
                serde_json::json!({
                    "type": "string",
                    "const": "notify_error"
                }),
            );

            let functions: HashMap<String, serde_json::Value> = tools_to_use
                .iter()
                .map(|tool| {
                    let func = tool.function.clone();

                    // Clone the existing parameters, which are expected to be a JSON object
                    let mut params = if let Value::Object(params) = &func.arguments {
                        params.clone()
                    } else {
                        Map::new()
                    };

                    // Insert the function's description at the top level, outside of properties
                    params.insert(
                        "description".to_string(),
                        Value::String(func.description.clone().unwrap_or_default()),
                    );

                    // Ensure 'properties' exists and is an object
                    let properties = params
                        .entry("properties".to_string())
                        .or_insert_with(|| json!({}))
                        .as_object_mut()
                        .unwrap();

                    // Insert the constant for the function name inside 'properties'
                    properties.insert(
                        "_name".to_string(),
                        json!({
                            "type": "string",
                            "const": func.name.clone(),
                            // "description": "The name of the function"
                        }),
                    );

                    // Check if 'required' exists, and it is an array. If not, create an empty array.
                    let required = params
                        .entry("required".to_string())
                        .or_insert_with(|| json!([]))
                        .as_array_mut()
                        .unwrap();

                    // Add 'name' to the 'required' array if it is not already present
                    if !required.iter().any(|r| r == "_name") {
                        required.push(json!("_name"));
                    }

                    (func.name, Value::Object(params))
                })
                .chain([(
                    "notify_error".to_string(),
                    serde_json::json!({
                        "properties": text_response_properties,
                        "required": ["error", "_name"],
                        "type": "object"
                    }),
                )])
                .collect();

            let tools = Tools {
                functions_map: FunctionsMap { functions },
                properties: Properties {
                    function: tools_to_use
                        .iter()
                        .map(|tool| FunctionRef {
                            ref_path: format!("#/$functions/{}", tool.function.name.clone()),
                        })
                        .chain(std::iter::once(FunctionRef {
                            ref_path: "#/$functions/notify_error".to_string(),
                        }))
                        .collect(),
                },
            };

            return Ok(Some(tools));
        }
        // Err(InferError::ToolError("No tools provided".to_string()))
        Ok(None)
    }
}

505
506
507
508
/// Batching logic
/// Will be launched in a background Tokio task
///
/// Batches requests and sends them to the inference server
509
#[allow(clippy::too_many_arguments)]
510
511
async fn batching_task(
    mut client: ShardedClient,
512
    waiting_served_ratio: f32,
513
    max_batch_prefill_tokens: u32,
514
    max_batch_total_tokens: u32,
515
    max_waiting_tokens: usize,
516
    max_batch_size: Option<usize>,
517
    queue: Queue,
518
    shared: Arc<Shared>,
519
    generation_health: Arc<AtomicBool>,
520
521
522
523
524
525
) {
    // Infinite loop
    loop {
        // Wait for a notification from the Infer struct
        shared.batching_task.notified().await;

526
        // Get the next batch from the queue
527
        // This batch might be smaller than the maximum batch size if there are not enough requests
528
        // waiting in the queue
529
        while let Some((mut entries, batch, span)) = queue
530
531
532
533
534
535
            .next_batch(
                None,
                max_batch_size,
                max_batch_prefill_tokens,
                max_batch_total_tokens,
            )
536
            .await
537
        {
538
            let mut cached_batch = prefill(&mut client, batch, &mut entries, &generation_health)
539
540
                .instrument(span)
                .await;
541
542
543
544
545
546
547
            let mut waiting_tokens = 1;

            // We loop until we do not receive any cached batch from the inference server (== until
            // all requests have met their stopping criteria)
            while let Some(batch) = cached_batch {
                // Get current batch info
                let batch_size = batch.size;
548
                let batch_max_tokens = batch.max_tokens;
549
                let mut batches = vec![batch];
550
                metrics::gauge!("tgi_batch_current_size", batch_size as f64);
551
552
553
554
555
556
557
558
559
560
561
                metrics::gauge!("tgi_batch_current_max_tokens", batch_max_tokens as f64);

                let min_size = if waiting_tokens >= max_waiting_tokens {
                    // If we didn't onboard any new requests since >= max_waiting_tokens, we try
                    // to add a new batch even though its size might be small
                    None
                } else {
                    // Minimum batch size
                    Some((batch_size as f32 * waiting_served_ratio).floor() as usize)
                };

562
                let token_budget = max_batch_total_tokens.saturating_sub(batch_max_tokens);
563
                let max_size = max_batch_size.map(|max_size| max_size - batch_size as usize);
564
565

                // Try to get a new batch
566
                if let Some((mut new_entries, new_batch, span)) = queue
567
                    .next_batch(min_size, max_size, max_batch_prefill_tokens, token_budget)
568
                    .await
569
570
571
572
573
574
575
                {
                    // Tracking metrics
                    if min_size.is_some() {
                        metrics::increment_counter!("tgi_batch_concat", "reason" => "backpressure");
                    } else {
                        metrics::increment_counter!("tgi_batch_concat", "reason" => "wait_exceeded");
                    }
576

577
578
579
580
581
582
583
584
585
586
587
588
                    entries.iter_mut().for_each(|(_, entry)| {
                        // Create a new span to add the info that this entry is waiting
                        // because a new batch is being computed
                        let entry_waiting_span = info_span!(parent: &entry.span, "waiting");
                        // Add relationships
                        span.follows_from(&entry_waiting_span);
                        entry_waiting_span.follows_from(&span);
                        // Update entry
                        entry.temp_span = Some(entry_waiting_span);
                    });

                    // Generate one token for this new batch to have the attention past in cache
589
590
591
592
                    let new_cached_batch =
                        prefill(&mut client, new_batch, &mut new_entries, &generation_health)
                            .instrument(span)
                            .await;
593
594
595
596
597
598
                    // Reset waiting counter
                    waiting_tokens = 1;
                    // Extend current batch with the new batch
                    if let Some(new_cached_batch) = new_cached_batch {
                        entries.extend(new_entries);
                        batches.push(new_cached_batch);
599
600
                    }
                }
601

602
603
604
605
606
607
                // Create span for this batch to add context to inference calls
                let next_batch_size = entries.len();
                let next_batch_span =
                    info_span!(parent: None, "batch", batch_size = next_batch_size);
                entries.iter_mut().for_each(|(_, entry)| {
                    // Create a new span to link the batch back to this entry
608
                    let entry_batch_span = info_span!(parent: &entry.span, "infer");
609
610
                    // Add relationships
                    next_batch_span.follows_from(&entry_batch_span);
611
612
613
614
                    entry_batch_span.follows_from(&next_batch_span);
                    // Update entry
                    entry.temp_span = Some(entry_batch_span);
                });
615

616
                cached_batch = decode(&mut client, batches, &mut entries, &generation_health)
617
618
                    .instrument(next_batch_span)
                    .await;
619
620
                waiting_tokens += 1;
            }
621
            metrics::gauge!("tgi_batch_current_size", 0.0);
622
            metrics::gauge!("tgi_batch_current_max_tokens", 0.0);
623
624
625
626
        }
    }
}

627
#[instrument(skip_all)]
628
629
630
async fn prefill(
    client: &mut ShardedClient,
    batch: Batch,
631
    entries: &mut IntMap<u64, Entry>,
632
    generation_health: &Arc<AtomicBool>,
633
) -> Option<CachedBatch> {
634
    let start_time = Instant::now();
635
    let batch_id = batch.id;
636
    metrics::increment_counter!("tgi_batch_inference_count", "method" => "prefill");
637
638

    match client.prefill(batch).await {
639
        Ok((generations, next_batch, timings)) => {
640
641
            // Update health
            generation_health.store(true, Ordering::SeqCst);
642
643

            let start_filtering_time = Instant::now();
644
            // Send generated tokens and filter stopped entries
645
646
647
            filter_send_generations(generations, entries);

            // Filter next batch and remove requests that were stopped
648
            let next_batch = filter_batch(client, next_batch, entries).await;
649

650
651
652
            metrics::histogram!("tgi_batch_forward_duration", timings.forward.as_secs_f64(), "method" => "prefill");
            metrics::histogram!("tgi_batch_decode_duration", timings.decode.as_secs_f64(), "method" => "prefill");
            metrics::histogram!("tgi_batch_filter_duration", start_filtering_time.elapsed().as_secs_f64(), "method" => "prefill");
653
            metrics::histogram!("tgi_batch_inference_duration", start_time.elapsed().as_secs_f64(), "method" => "prefill");
654
655
656
657
658
            metrics::increment_counter!("tgi_batch_inference_success", "method" => "prefill");
            next_batch
        }
        // If we have an error, we discard the whole batch
        Err(err) => {
659
660
            // Update health
            generation_health.store(false, Ordering::SeqCst);
661
            let _ = client.clear_cache(Some(batch_id)).await;
662
663
664
665
666
667
668
669
670
671
            send_errors(err, entries);
            metrics::increment_counter!("tgi_batch_inference_failure", "method" => "prefill");
            None
        }
    }
}

#[instrument(skip_all)]
async fn decode(
    client: &mut ShardedClient,
672
    batches: Vec<CachedBatch>,
673
    entries: &mut IntMap<u64, Entry>,
674
    generation_health: &Arc<AtomicBool>,
675
) -> Option<CachedBatch> {
676
    let start_time = Instant::now();
677
    let batch_ids: Vec<u64> = batches.iter().map(|b| b.id).collect();
678
    metrics::increment_counter!("tgi_batch_inference_count", "method" => "decode");
679
680

    match client.decode(batches).await {
681
        Ok((generations, next_batch, timings)) => {
682
683
            // Update health
            generation_health.store(true, Ordering::SeqCst);
684
685

            let start_filtering_time = Instant::now();
686
            // Send generated tokens and filter stopped entries
687
688
689
            filter_send_generations(generations, entries);

            // Filter next batch and remove requests that were stopped
690
            let next_batch = filter_batch(client, next_batch, entries).await;
691

692
693
694
695
696
697
            if let Some(concat_duration) = timings.concat {
                metrics::histogram!("tgi_batch_concat_duration", concat_duration.as_secs_f64(), "method" => "decode");
            }
            metrics::histogram!("tgi_batch_forward_duration", timings.forward.as_secs_f64(), "method" => "decode");
            metrics::histogram!("tgi_batch_decode_duration", timings.decode.as_secs_f64(), "method" => "decode");
            metrics::histogram!("tgi_batch_filter_duration", start_filtering_time.elapsed().as_secs_f64(), "method" => "decode");
698
            metrics::histogram!("tgi_batch_inference_duration", start_time.elapsed().as_secs_f64(), "method" => "decode");
699
            metrics::increment_counter!("tgi_batch_inference_success", "method" => "decode");
700
701
702
703
            next_batch
        }
        // If we have an error, we discard the whole batch
        Err(err) => {
704
            generation_health.store(false, Ordering::SeqCst);
705
706
707
            for id in batch_ids {
                let _ = client.clear_cache(Some(id)).await;
            }
708
            send_errors(err, entries);
709
            metrics::increment_counter!("tgi_batch_inference_failure", "method" => "decode");
710
711
712
713
714
            None
        }
    }
}

715
716
/// Filter a `batch` and remove all requests not present in `entries`
#[instrument(skip_all)]
717
718
async fn filter_batch(
    client: &mut ShardedClient,
719
    next_batch: Option<CachedBatch>,
720
    entries: &IntMap<u64, Entry>,
721
) -> Option<CachedBatch> {
722
723
724
725
726
727
728
729
730
731
    let mut batch = next_batch?;

    // No need to filter
    if batch.size as usize == entries.len() {
        return Some(batch);
    }

    let id = batch.id;

    // Retain only requests that are still in entries
732
    batch.request_ids.retain(|id| entries.contains_key(id));
733

734
    if batch.request_ids.is_empty() {
735
736
737
738
739
740
741
742
743
        // All requests have been filtered out
        // Next batch is now empty
        // Clear it from the Python shards cache
        // We unwrap here as we need to panic since we cannot recover if this method fails
        client.clear_cache(Some(id)).await.unwrap();
        None
    } else {
        // Filter Python shard cache
        // We unwrap here as we need to panic since we cannot recover if this method fails
744
        client.filter_batch(id, batch.request_ids).await.unwrap()
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
    }
}

/// Send one or multiple `InferStreamResponse` to Infer for all `entries`
/// and filter entries
#[instrument(skip_all)]
fn filter_send_generations(generations: Vec<Generation>, entries: &mut IntMap<u64, Entry>) {
    generations.into_iter().for_each(|generation| {
        let id = generation.request_id;
        // Get entry
        // We can `expect` here as the request id should always be in the entries
        let entry = entries
            .get(&id)
            .expect("ID not found in entries. This is a bug.");

        // Create and enter a span to link this function back to the entry
        let _span = info_span!(parent: entry.temp_span.as_ref().expect("batch_span is None. This is a bug."), "send_generation", generation = ?generation).entered();
        // Send generation responses back to the infer task
        // If the receive an error from the Flume channel, it means that the client dropped the
        // request and we need to stop generating hence why we unwrap_or(true)
        let stopped = send_responses(generation, entry).map_err(|err| {
OlivierDehaene's avatar
OlivierDehaene committed
766
            tracing::error!("Entry response channel error.");
767
768
769
770
771
772
773
774
775
776
777
778
779
            metrics::increment_counter!("tgi_request_failure", "err" => "dropped");
            err
        }).unwrap_or(true);
        if stopped {
            entries.remove(&id).expect("ID not found in entries. This is a bug.");
        }
    });
}

/// Send responses through the `entry` response channel
fn send_responses(
    generation: Generation,
    entry: &Entry,
OlivierDehaene's avatar
OlivierDehaene committed
780
) -> Result<bool, Box<SendError<Result<InferStreamResponse, InferError>>>> {
781
    // Return directly if the channel is disconnected
OlivierDehaene's avatar
OlivierDehaene committed
782
783
    if entry.response_tx.is_closed() {
        metrics::increment_counter!("tgi_request_failure", "err" => "dropped");
784
785
786
        return Ok(true);
    }

787
788
789
790
    let mut stopped = false;

    if let Some(prefill_tokens) = generation.prefill_tokens {
        // Send message
OlivierDehaene's avatar
OlivierDehaene committed
791
792
793
        entry
            .response_tx
            .send(Ok(InferStreamResponse::Prefill(prefill_tokens)))?;
794
795
796
    }

    // Create last Token
Nicolas Patry's avatar
Nicolas Patry committed
797
798
799
800
801
802
    let tokens_ = generation.tokens.expect("Non empty tokens in generation");
    let n = tokens_.ids.len();
    metrics::histogram!("tgi_request_skipped_tokens", (n - 1) as f64);
    let mut iterator = tokens_
        .ids
        .into_iter()
803
804
805
        .zip(tokens_.logprobs)
        .zip(tokens_.texts)
        .zip(tokens_.is_special)
Nicolas Patry's avatar
Nicolas Patry committed
806
807
808
809
810
811
812
813
814
815
        .enumerate()
        .peekable();
    while let Some((i, (((id, logprob), text), special))) = iterator.next() {
        let token = Token {
            id,
            text,
            logprob,
            special,
        };
        let top_tokens = if let Some(top_tokens_) = generation.top_tokens.get(i) {
Nicolas Patry's avatar
Nicolas Patry committed
816
817
            top_tokens_
                .ids
Nicolas Patry's avatar
Nicolas Patry committed
818
819
820
821
822
                .iter()
                .zip(top_tokens_.logprobs.iter())
                .zip(top_tokens_.texts.iter())
                .zip(top_tokens_.is_special.iter())
                .map(|(((&id, &logprob), text), &special)| Token {
Nicolas Patry's avatar
Nicolas Patry committed
823
                    id,
Nicolas Patry's avatar
Nicolas Patry committed
824
                    text: text.to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
825
826
                    logprob,
                    special,
Nicolas Patry's avatar
Nicolas Patry committed
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
                })
                .collect()
        } else {
            vec![]
        };
        match (&generation.generated_text, iterator.peek()) {
            (Some(generated_text), None) => {
                // Generation has ended
                stopped = true;
                // Send message
                entry.response_tx.send(Ok(InferStreamResponse::End {
                    token,
                    top_tokens,
                    generated_text: generated_text.clone(),
                    queued: entry.queue_time,
                    start: entry.batch_time.unwrap(),
                }))?;
            }
            _ => {
                // Send message
                entry
                    .response_tx
                    .send(Ok(InferStreamResponse::Intermediate { token, top_tokens }))?;
            }
        }
Nicolas Patry's avatar
Nicolas Patry committed
852
853
    }

854
855
856
    Ok(stopped)
}

857
/// Send errors to Infer for all `entries`
858
859
#[instrument(skip_all)]
fn send_errors(error: ClientError, entries: &mut IntMap<u64, Entry>) {
860
    entries.drain().for_each(|(_, entry)| {
861
862
863
        // Create and enter a span to link this function back to the entry
        let _send_error_span = info_span!(parent: entry.temp_span.as_ref().expect("batch_span is None. This is a bug."), "send_error").entered();
        let err = InferError::GenerationError(error.to_string());
864
        metrics::increment_counter!("tgi_request_failure", "err" => "generation");
865
866
        tracing::error!("{err}");

867
868
869
        // unwrap_or is valid here as we don't care if the receiver is gone.
        entry
            .response_tx
OlivierDehaene's avatar
OlivierDehaene committed
870
            .send(Err(err))
871
872
873
874
875
876
877
            .unwrap_or(());
    });
}

#[derive(Debug)]
pub(crate) enum InferStreamResponse {
    // Optional first message
Nicolas Patry's avatar
Nicolas Patry committed
878
    Prefill(Tokens),
879
    // Intermediate messages
Nicolas Patry's avatar
Nicolas Patry committed
880
881
882
883
    Intermediate {
        token: Token,
        top_tokens: Vec<Token>,
    },
884
885
886
    // Last message
    End {
        token: Token,
Nicolas Patry's avatar
Nicolas Patry committed
887
        top_tokens: Vec<Token>,
888
889
890
891
892
893
894
895
        generated_text: GeneratedText,
        start: Instant,
        queued: Instant,
    },
}

#[derive(Debug)]
pub(crate) struct InferResponse {
896
897
898
899
    /// input_length is the input as perceived by the rust tokenizer in the
    /// validation pathway. It is redundant with prefill.len() but prefill
    /// has data only if the user asked for it. This will always be filled.
    pub(crate) _input_length: u32,
900
    pub(crate) prefill: Vec<PrefillToken>,
901
902
903
904
    pub(crate) tokens: Vec<Token>,
    pub(crate) generated_text: GeneratedText,
    pub(crate) queued: Instant,
    pub(crate) start: Instant,
Nicolas Patry's avatar
Nicolas Patry committed
905
    pub(crate) top_tokens: Vec<Vec<Token>>,
906
907
908
909
910
911
912
913
914
915
916
917
}

#[derive(Debug, Error)]
pub enum InferError {
    #[error("Request failed during generation: {0}")]
    GenerationError(String),
    #[error("Model is overloaded")]
    Overloaded(#[from] TryAcquireError),
    #[error("Input validation error: {0}")]
    ValidationError(#[from] ValidationError),
    #[error("Incomplete generation")]
    IncompleteGeneration,
918
919
    #[error("Template error: {0}")]
    TemplateError(#[from] minijinja::Error),
920
921
    #[error("Tool error: {0}")]
    ToolError(String),
922
}
923
924
925
926
927
928
929
930

impl InferError {
    pub(crate) fn error_type(&self) -> &str {
        match self {
            InferError::GenerationError(_) => "generation",
            InferError::Overloaded(_) => "overloaded",
            InferError::ValidationError(_) => "validation",
            InferError::IncompleteGeneration => "incomplete_generation",
931
            InferError::TemplateError(_) => "template_error",
932
            InferError::ToolError(_) => "tool_error",
933
934
935
        }
    }
}
936
937
938
939
940

// tests
#[cfg(test)]
mod tests {
    use crate::infer::raise_exception;
Nicolas Patry's avatar
Nicolas Patry committed
941
    use crate::{ChatTemplateInputs, TextMessage};
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
    use minijinja::Environment;

    #[test]
    fn test_chat_template() {
        let env = Environment::new();

        let source = r#"
        {% for message in messages %}
            {% if message['role'] == 'system' %}
                {% if message['content']%}
                    {{'### System:\n' + message['content']+'\n\n'}}
                {% endif %}
            {% elif message['role'] == 'user' %}
                {{'### User:\n' + message['content']+'\n\n'}}
            {% elif message['role'] == 'assistant' %}
                {{'### Assistant:\n'  + message['content']}}
            {% endif %}
            {% if loop.last and add_generation_prompt %}
                {{ '### Assistant:\n' }}
            {% endif %}
        {% endfor %}"#;

        // trim all the whitespace
        let source = source
            .lines()
            .map(|line| line.trim())
            .collect::<Vec<&str>>()
            .join("");

        let tmpl = env.template_from_str(&source);

        let chat_template_inputs = ChatTemplateInputs {
            messages: vec![
Nicolas Patry's avatar
Nicolas Patry committed
975
                TextMessage {
976
                    role: "user".to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
977
                    content: "Hi!".to_string(),
978
                },
Nicolas Patry's avatar
Nicolas Patry committed
979
                TextMessage {
980
                    role: "assistant".to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
981
                    content: "Hello how can I help?".to_string(),
982
                },
Nicolas Patry's avatar
Nicolas Patry committed
983
                TextMessage {
984
                    role: "user".to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
985
                    content: "What is Deep Learning?".to_string(),
986
                },
Nicolas Patry's avatar
Nicolas Patry committed
987
                TextMessage {
988
                    role: "assistant".to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
989
                    content: "magic!".to_string(),
990
991
992
993
                },
            ],
            bos_token: Some("[BOS]"),
            eos_token: Some("[EOS]"),
994
            add_generation_prompt: true,
995
            ..Default::default()
996
997
998
999
1000
1001
        };

        let result = tmpl.unwrap().render(chat_template_inputs).unwrap();

        assert_eq!(
            result,
1002
            "### User:\nHi!\n\n### Assistant:\nHello how can I help?### User:\nWhat is Deep Learning?\n\n### Assistant:\nmagic!### Assistant:\n"
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
        );
    }

    #[test]
    fn test_chat_template_invalid_with_raise() {
        let mut env = Environment::new();
        env.add_function("raise_exception", raise_exception);

        let source = r#"
        {{ bos_token }}
        {% for message in messages %}
        {% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}
        {{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}
        {% endif %}
        {% if message['role'] == 'user' %}
        {{ '[INST] ' + message['content'] + ' [/INST]' }}
        {% elif message['role'] == 'assistant' %}
        {{ message['content'] + eos_token}}
        {% else %}
        {{ raise_exception('Only user and assistant roles are supported!') }}
        {% endif %}
        {% endfor %}"#;

        // trim all the whitespace
        let source = source
            .lines()
            .map(|line| line.trim())
            .collect::<Vec<&str>>()
            .join("");

        let tmpl = env.template_from_str(&source);

        let chat_template_inputs = ChatTemplateInputs {
            messages: vec![
Nicolas Patry's avatar
Nicolas Patry committed
1037
                TextMessage {
1038
                    role: "user".to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
1039
                    content: "Hi!".to_string(),
1040
                },
Nicolas Patry's avatar
Nicolas Patry committed
1041
                TextMessage {
1042
                    role: "user".to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
1043
                    content: "Hi again!".to_string(),
1044
                },
Nicolas Patry's avatar
Nicolas Patry committed
1045
                TextMessage {
1046
                    role: "assistant".to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
1047
                    content: "Hello how can I help?".to_string(),
1048
                },
Nicolas Patry's avatar
Nicolas Patry committed
1049
                TextMessage {
1050
                    role: "user".to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
1051
                    content: "What is Deep Learning?".to_string(),
1052
                },
Nicolas Patry's avatar
Nicolas Patry committed
1053
                TextMessage {
1054
                    role: "assistant".to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
1055
                    content: "magic!".to_string(),
1056
1057
1058
1059
                },
            ],
            bos_token: Some("[BOS]"),
            eos_token: Some("[EOS]"),
1060
            add_generation_prompt: true,
1061
            ..Default::default()
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
        };

        let result = tmpl.unwrap().render(chat_template_inputs); //.err().unwrap();

        match result {
            Ok(_) => panic!("Should have failed"),
            Err(e) => {
                assert_eq!(
                    e.detail().unwrap(),
                    "Conversation roles must alternate user/assistant/user/assistant/..."
                );
            }
        }
    }

    #[test]
    fn test_chat_template_valid_with_raise() {
        let mut env = Environment::new();
        env.add_function("raise_exception", raise_exception);

        let source = r#"
        {{ bos_token }}
        {% for message in messages %}
        {% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}
        {{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}
        {% endif %}
        {% if message['role'] == 'user' %}
        {{ '[INST] ' + message['content'] + ' [/INST]' }}
        {% elif message['role'] == 'assistant' %}
        {{ message['content'] + eos_token}}
        {% else %}
        {{ raise_exception('Only user and assistant roles are supported!') }}
        {% endif %}
        {% endfor %}"#;

        // trim all the whitespace
        let source = source
            .lines()
            .map(|line| line.trim())
            .collect::<Vec<&str>>()
            .join("");

        let tmpl = env.template_from_str(&source);

        let chat_template_inputs = ChatTemplateInputs {
            messages: vec![
Nicolas Patry's avatar
Nicolas Patry committed
1108
                TextMessage {
1109
                    role: "user".to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
1110
                    content: "Hi!".to_string(),
1111
                },
Nicolas Patry's avatar
Nicolas Patry committed
1112
                TextMessage {
1113
                    role: "assistant".to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
1114
                    content: "Hello how can I help?".to_string(),
1115
                },
Nicolas Patry's avatar
Nicolas Patry committed
1116
                TextMessage {
1117
                    role: "user".to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
1118
                    content: "What is Deep Learning?".to_string(),
1119
                },
Nicolas Patry's avatar
Nicolas Patry committed
1120
                TextMessage {
1121
                    role: "assistant".to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
1122
                    content: "magic!".to_string(),
1123
1124
1125
1126
                },
            ],
            bos_token: Some("[BOS]"),
            eos_token: Some("[EOS]"),
1127
            add_generation_prompt: true,
1128
            ..Default::default()
1129
1130
1131
1132
1133
        };

        let result = tmpl.unwrap().render(chat_template_inputs).unwrap();
        assert_eq!(result, "[BOS][INST] Hi! [/INST]Hello how can I help?[EOS][INST] What is Deep Learning? [/INST]magic![EOS]");
    }
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158

    #[test]
    fn test_chat_template_valid_with_add_generation_prompt() {
        let mut env = Environment::new();
        env.add_function("raise_exception", raise_exception);

        let source = r#"
        {% for message in messages %}
        {{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}
        {% endfor %}
        {% if add_generation_prompt %}
            {{ '<|im_start|>assistant\n' }}
        {% endif %}"#;

        // trim all the whitespace
        let source = source
            .lines()
            .map(|line| line.trim())
            .collect::<Vec<&str>>()
            .join("");

        let tmpl = env.template_from_str(&source);

        let chat_template_inputs = ChatTemplateInputs {
            messages: vec![
Nicolas Patry's avatar
Nicolas Patry committed
1159
                TextMessage {
1160
                    role: "user".to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
1161
                    content: "Hi!".to_string(),
1162
                },
Nicolas Patry's avatar
Nicolas Patry committed
1163
                TextMessage {
1164
                    role: "assistant".to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
1165
                    content: "Hello how can I help?".to_string(),
1166
                },
Nicolas Patry's avatar
Nicolas Patry committed
1167
                TextMessage {
1168
                    role: "user".to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
1169
                    content: "What is Deep Learning?".to_string(),
1170
                },
Nicolas Patry's avatar
Nicolas Patry committed
1171
                TextMessage {
1172
                    role: "assistant".to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
1173
                    content: "magic!".to_string(),
1174
1175
1176
1177
1178
                },
            ],
            bos_token: Some("[BOS]"),
            eos_token: Some("[EOS]"),
            add_generation_prompt: true,
1179
            ..Default::default()
1180
1181
1182
1183
1184
        };

        let result = tmpl.unwrap().render(chat_template_inputs).unwrap();
        assert_eq!(result, "<|im_start|>user\nHi!<|im_end|>\n<|im_start|>assistant\nHello how can I help?<|im_end|>\n<|im_start|>user\nWhat is Deep Learning?<|im_end|>\n<|im_start|>assistant\nmagic!<|im_end|>\n<|im_start|>assistant\n");
    }
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195

    struct ChatTemplateTestItem {
        name: &'static str,
        chat_template: &'static str,
        input: ChatTemplateInputs<'static>,
        target: &'static str,
    }

    #[test]
    fn test_many_chat_templates() {
        let example_chat = vec![
Nicolas Patry's avatar
Nicolas Patry committed
1196
            TextMessage {
1197
                role: "user".to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
1198
                content: "Hello, how are you?".to_string(),
1199
            },
Nicolas Patry's avatar
Nicolas Patry committed
1200
            TextMessage {
1201
                role: "assistant".to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
1202
                content: "I'm doing great. How can I help you today?".to_string(),
1203
            },
Nicolas Patry's avatar
Nicolas Patry committed
1204
            TextMessage {
1205
                role: "user".to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
1206
                content: "I'd like to show off how chat templating works!".to_string(),
1207
1208
1209
            },
        ];

Nicolas Patry's avatar
Nicolas Patry committed
1210
        let example_chat_with_system = [TextMessage {
1211
            role: "system".to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
1212
1213
            content: "You are a friendly chatbot who always responds in the style of a pirate"
                .to_string(),
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
        }]
        .iter()
        .chain(&example_chat)
        .cloned()
        .collect::<Vec<_>>();

        let test_default_templates = vec![
            ChatTemplateTestItem {
                name: "_base",
                chat_template: "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\\n' + message['content'] + '<|im_end|>' + '\\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\\n' }}{% endif %}",
1224
                input: ChatTemplateInputs {
1225
1226
1227
1228
                    messages: example_chat.clone(),
                    add_generation_prompt: false,
                    bos_token: Some(""),
                    eos_token: Some(""),
1229
                    ..Default::default()
1230
1231
1232
1233
1234
1235
                },
                target: "<|im_start|>user\nHello, how are you?<|im_end|>\n<|im_start|>assistant\nI'm doing great. How can I help you today?<|im_end|>\n<|im_start|>user\nI'd like to show off how chat templating works!<|im_end|>\n",
            },
            ChatTemplateTestItem {
                name: "blenderbot",
                chat_template: "{% for message in messages %}{% if message['role'] == 'user' %}{{ ' ' }}{% endif %}{{ message['content'] }}{% if not loop.last %}{{ '  ' }}{% endif %}{% endfor %}{{ eos_token }}",
1236
                input: ChatTemplateInputs {
1237
1238
1239
1240
                    messages: example_chat.clone(),
                    add_generation_prompt: false,
                    bos_token: Some(""),
                    eos_token: Some("</s>"),
1241
                    ..Default::default()
1242
1243
1244
1245
1246
1247
                },
                target: " Hello, how are you?  I'm doing great. How can I help you today?   I'd like to show off how chat templating works!</s>",
            },
            ChatTemplateTestItem {
                name: "blenderbot_small",
                chat_template: "{% for message in messages %}{% if message['role'] == 'user' %}{{ ' ' }}{% endif %}{{ message['content'] }}{% if not loop.last %}{{ '  ' }}{% endif %}{% endfor %}{{ eos_token }}",
1248
                input: ChatTemplateInputs {
1249
1250
1251
1252
                    messages: example_chat.clone(),
                    add_generation_prompt: false,
                    bos_token: Some(""),
                    eos_token: Some("</s>"),
1253
                    ..Default::default()
1254
1255
1256
1257
1258
1259
                },
                target: " Hello, how are you?  I'm doing great. How can I help you today?   I'd like to show off how chat templating works!</s>",
            },
            ChatTemplateTestItem {
                name: "bloom",
                chat_template: "{% for message in messages %}{{ message.content }}{{ eos_token }}{% endfor %}",
1260
                input: ChatTemplateInputs {
1261
1262
1263
1264
                    messages: example_chat.clone(),
                    add_generation_prompt: false,
                    bos_token: Some(""),
                    eos_token: Some("</s>"),
1265
                    ..Default::default()
1266
1267
1268
1269
1270
1271
                },
                target: "Hello, how are you?</s>I'm doing great. How can I help you today?</s>I'd like to show off how chat templating works!</s>",
            },
            ChatTemplateTestItem {
                name: "gpt_neox",
                chat_template: "{% for message in messages %}{{ message.content }}{{ eos_token }}{% endfor %}",
1272
                input: ChatTemplateInputs {
1273
1274
1275
1276
                    messages: example_chat.clone(),
                    add_generation_prompt: false,
                    bos_token: Some(""),
                    eos_token: Some("<|endoftext|>"),
1277
                    ..Default::default()
1278
1279
1280
1281
1282
1283
                },
                target: "Hello, how are you?<|endoftext|>I'm doing great. How can I help you today?<|endoftext|>I'd like to show off how chat templating works!<|endoftext|>",
            },
            ChatTemplateTestItem {
                name: "gpt2",
                chat_template: "{% for message in messages %}{{ message.content }}{{ eos_token }}{% endfor %}",
1284
1285
1286
1287
1288
                input: ChatTemplateInputs {
                    messages: example_chat.clone(),
                    add_generation_prompt: false,
                    bos_token: Some(""),
                    eos_token: Some("<|endoftext|>"),
1289
                    ..Default::default()
1290
                },
1291
                target: "Hello, how are you?<|endoftext|>I'm doing great. How can I help you today?<|endoftext|>I'd like to show off how chat templating works!<|endoftext|>",
1292
1293
1294
1295
1296
            },
            ChatTemplateTestItem {
                name: "llama",
                // NOTE: the `.strip()` has been replaced with `| trim` in the following template
                chat_template: "{% if messages[0]['role'] == 'system' %}{% set loop_messages = messages[1:] %}{% set system_message = messages[0]['content'] %}{% elif USE_DEFAULT_PROMPT == true and not '<<SYS>>' in messages[0]['content'] %}{% set loop_messages = messages %}{% set system_message = 'DEFAULT_SYSTEM_MESSAGE' %}{% else %}{% set loop_messages = messages %}{% set system_message = false %}{% endif %}{% for message in loop_messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if loop.index0 == 0 and system_message != false %}{% set content = '<<SYS>>\\n' + system_message + '\\n<</SYS>>\\n\\n' + message['content'] %}{% else %}{% set content = message['content'] %}{% endif %}{% if message['role'] == 'user' %}{{ bos_token +'[INST] ' + content | trim + ' [/INST]' }}{% elif message['role'] == 'system' %}{{ '<<SYS>>\\n' + content | trim + '\\n<</SYS>>\\n\\n' }}{% elif message['role'] == 'assistant' %}{{ ' ' + content | trim + ' ' + eos_token }}{% endif %}{% endfor %}",
1297
1298
1299
1300
1301
                input: ChatTemplateInputs {
                    messages: example_chat_with_system.clone(),
                    add_generation_prompt: true,
                    bos_token: Some("<s>"),
                    eos_token: Some("</s>"),
1302
                    ..Default::default()
1303
                },
1304
                target: "<s>[INST] <<SYS>>\nYou are a friendly chatbot who always responds in the style of a pirate\n<</SYS>>\n\nHello, how are you? [/INST] I'm doing great. How can I help you today? </s><s>[INST] I'd like to show off how chat templating works! [/INST]",
1305
1306
1307
1308
            },
            ChatTemplateTestItem {
                name: "whisper",
                chat_template: "{% for message in messages %}{{ message.content }}{{ eos_token }}{% endfor %}",
1309
1310
1311
1312
1313
                input: ChatTemplateInputs {
                    messages: example_chat.clone(),
                    add_generation_prompt: true,
                    bos_token: Some(""),
                    eos_token: Some("<|endoftext|>"),
1314
                    ..Default::default()
1315
                },
1316
1317
                target: "Hello, how are you?<|endoftext|>I'm doing great. How can I help you today?<|endoftext|>I'd like to show off how chat templating works!<|endoftext|>",
            },
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
        ];

        #[allow(unused_variables)] // name is unused
        for ChatTemplateTestItem {
            name,
            chat_template,
            input,
            target,
        } in test_default_templates
        {
            let mut env = Environment::new();
            env.add_function("raise_exception", raise_exception);
1330
            let tmpl = env.template_from_str(chat_template);
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
            let result = tmpl.unwrap().render(input).unwrap();
            assert_eq!(result, target);
        }

        let test_custom_templates = vec![
            ChatTemplateTestItem {
                name: "HuggingFaceH4/zephyr-7b-beta (add_generation_prompt=false)",
                chat_template: "{% for message in messages %}\n{% if message['role'] == 'user' %}\n{{ '<|user|>\\n' + message['content'] + eos_token }}\n{% elif message['role'] == 'system' %}\n{{ '<|system|>\\n' + message['content'] + eos_token }}\n{% elif message['role'] == 'assistant' %}\n{{ '<|assistant|>\\n'  + message['content'] + eos_token }}\n{% endif %}\n{% if loop.last and add_generation_prompt %}\n{{ '<|assistant|>' }}\n{% endif %}\n{% endfor %}",
                input: ChatTemplateInputs {
                    messages: example_chat_with_system.clone(),
                    add_generation_prompt: false,
                    bos_token: Some(""),
1343
                    eos_token: Some("</s>"),
1344
                    ..Default::default()
1345
1346
1347
1348
1349
1350
1351
1352
                },
                target: "<|system|>\nYou are a friendly chatbot who always responds in the style of a pirate</s><|user|>\nHello, how are you?</s><|assistant|>\nI'm doing great. How can I help you today?</s><|user|>\nI'd like to show off how chat templating works!</s>",
            },
            ChatTemplateTestItem {
                name: "HuggingFaceH4/zephyr-7b-beta (add_generation_prompt=true)",
                chat_template: "{% for message in messages %}\n{% if message['role'] == 'user' %}\n{{ '<|user|>\\n' + message['content'] + eos_token }}\n{% elif message['role'] == 'system' %}\n{{ '<|system|>\\n' + message['content'] + eos_token }}\n{% elif message['role'] == 'assistant' %}\n{{ '<|assistant|>\\n'  + message['content'] + eos_token }}\n{% endif %}\n{% if loop.last and add_generation_prompt %}\n{{ '<|assistant|>' }}\n{% endif %}\n{% endfor %}",
                input: ChatTemplateInputs {
                    messages: vec![
Nicolas Patry's avatar
Nicolas Patry committed
1353
                        TextMessage{
1354
                            role: "system".to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
1355
                            content: "You are a friendly chatbot who always responds in the style of a pirate".to_string(),
1356
                        },
Nicolas Patry's avatar
Nicolas Patry committed
1357
                        TextMessage{
1358
                            role: "user".to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
1359
                            content: "How many helicopters can a human eat in one sitting?".to_string(),
1360
1361
1362
1363
1364
                        },
                    ],
                    add_generation_prompt: true,
                    bos_token: Some(""),
                    eos_token: Some("</s>"),
1365
                    ..Default::default()
1366
                },
1367
                target: "<|system|>\nYou are a friendly chatbot who always responds in the style of a pirate</s><|user|>\nHow many helicopters can a human eat in one sitting?</s><|assistant|>",
1368
1369
1370
1371
1372
1373
1374
1375
1376
            },
            ChatTemplateTestItem {
                name: "HuggingFaceH4/zephyr-7b-gemma-v0.1",
                chat_template: "{% if messages[0]['role'] == 'user' or messages[0]['role'] == 'system' %}{{ bos_token }}{% endif %}{% for message in messages %}{{ '<|im_start|>' + message['role'] + '\\n' + message['content'] + '<|im_end|>' + '\\n' }}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% elif messages[-1]['role'] == 'assistant' %}{{ eos_token }}{% endif %}",
                input: ChatTemplateInputs {
                    messages: example_chat.clone(),
                    add_generation_prompt: false,
                    bos_token: Some("<bos>"),
                    eos_token: Some("<eos>"),
1377
                    ..Default::default()
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
                },
                target: "<bos><|im_start|>user\nHello, how are you?<|im_end|>\n<|im_start|>assistant\nI'm doing great. How can I help you today?<|im_end|>\n<|im_start|>user\nI'd like to show off how chat templating works!<|im_end|>\n",
            },
            ChatTemplateTestItem {
                name: "mistralai/Mistral-7B-Instruct-v0.1",
                chat_template: "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ '[INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ message['content'] + eos_token + ' ' }}{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}",
                input: ChatTemplateInputs {
                    messages: example_chat.clone(),
                    add_generation_prompt: false,
                    bos_token: Some("<s>"),
                    eos_token: Some("</s>"),
1389
                    ..Default::default()
1390
                },
1391
                target: "<s>[INST] Hello, how are you? [/INST]I'm doing great. How can I help you today?</s> [INST] I'd like to show off how chat templating works! [/INST]",
1392
1393
1394
1395
1396
1397
1398
1399
1400
            },
            ChatTemplateTestItem {
                name: "mistralai/Mixtral-8x7B-Instruct-v0.1",
                chat_template: "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ '[INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ message['content'] + eos_token}}{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}",
                input: ChatTemplateInputs {
                    messages: example_chat.clone(),
                    add_generation_prompt: false,
                    bos_token: Some("<s>"),
                    eos_token: Some("</s>"),
1401
                    ..Default::default()
1402
1403
1404
1405
1406
1407
1408
                },
                target: "<s>[INST] Hello, how are you? [/INST]I'm doing great. How can I help you today?</s>[INST] I'd like to show off how chat templating works! [/INST]",
            },
            ChatTemplateTestItem {
                name: "cognitivecomputations/dolphin-2.5-mixtral-8x7b",
                chat_template: "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% for message in messages %}{{'<|im_start|>' + message['role'] + '\\n' + message['content'] + '<|im_end|>' + '\\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\\n' }}{% endif %}",
                input: ChatTemplateInputs {
1409
                    messages: example_chat.clone(),
1410
1411
1412
                    add_generation_prompt: false,
                    bos_token: Some("<s>"),
                    eos_token: Some("</s>"),
1413
                    ..Default::default()
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
                },
                target: "<|im_start|>user\nHello, how are you?<|im_end|>\n<|im_start|>assistant\nI'm doing great. How can I help you today?<|im_end|>\n<|im_start|>user\nI'd like to show off how chat templating works!<|im_end|>\n",
            },
            ChatTemplateTestItem {
                name: "openchat/openchat-3.5-0106",
                // `.title()` has been replaced with `| upper` in the following template
                chat_template: "{{ bos_token }}{% for message in messages %}{{ 'GPT4 Correct ' + (message['role'] | title) + ': ' + message['content'] + '<|end_of_turn|>'}}{% endfor %}{% if add_generation_prompt %}{{ 'GPT4 Correct Assistant:' }}{% endif %}",
                input: ChatTemplateInputs {
                    messages: example_chat.clone(),
                    add_generation_prompt: false,
                    bos_token: Some("<s>"),
                    eos_token: Some("</s>"),
1426
                    ..Default::default()
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
                },
                target: "<s>GPT4 Correct User: Hello, how are you?<|end_of_turn|>GPT4 Correct Assistant: I'm doing great. How can I help you today?<|end_of_turn|>GPT4 Correct User: I'd like to show off how chat templating works!<|end_of_turn|>",
            },
            ChatTemplateTestItem {
                name: "upstage/SOLAR-10.7B-Instruct-v1.0",
                chat_template: "{% for message in messages %}{{ message.content }}{{ eos_token }}{% endfor %}",
                input: ChatTemplateInputs {
                    messages: example_chat.clone(),
                    add_generation_prompt: false,
                    bos_token: Some("<s>"),
                    eos_token: Some("</s>"),
1438
                    ..Default::default()
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
                },
                target: "Hello, how are you?</s>I'm doing great. How can I help you today?</s>I'd like to show off how chat templating works!</s>",
            },
            ChatTemplateTestItem {
                name: "codellama/CodeLlama-70b-Instruct-hf",
                // NOTE: `.strip()` has been replaced with `| trim` in the following template
                chat_template: "{% if messages[0]['role'] == 'system' %}{% set user_index = 1 %}{% else %}{% set user_index = 0 %}{% endif %}{% for message in messages %}{% if (message['role'] == 'user') != ((loop.index0 + user_index) % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if loop.index0 == 0 %}{{ '<s>' }}{% endif %}{% set content = 'Source: ' + message['role'] + '\\n\\n ' + message['content'] | trim %}{{ content + ' <step> ' }}{% endfor %}{{'Source: assistant\\nDestination: user\\n\\n '}}",
                input: ChatTemplateInputs {
                    messages: example_chat.clone(),
                    add_generation_prompt: false,
                    bos_token: Some("<s>"),
                    eos_token: Some("</s>"),
1451
                    ..Default::default()
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
                },
                target: "<s>Source: user\n\n Hello, how are you? <step> Source: assistant\n\n I'm doing great. How can I help you today? <step> Source: user\n\n I'd like to show off how chat templating works! <step> Source: assistant\nDestination: user\n\n ",
            },
            ChatTemplateTestItem {
                name: "Deci/DeciLM-7B-instruct",
                chat_template: "{% for message in messages %}\n{% if message['role'] == 'user' %}\n{{ '### User:\\n' + message['content'] }}\n{% elif message['role'] == 'system' %}\n{{ '### System:\\n' + message['content'] }}\n{% elif message['role'] == 'assistant' %}\n{{ '### Assistant:\\n'  + message['content'] }}\n{% endif %}\n{% if loop.last and add_generation_prompt %}\n{{ '### Assistant:' }}\n{% endif %}\n{% endfor %}",
                input: ChatTemplateInputs {
                    messages: example_chat.clone(),
                    add_generation_prompt: false,
                    bos_token: Some("<s>"),
                    eos_token: Some("</s>"),
1463
                    ..Default::default()
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
                },
                target: "### User:\nHello, how are you?### Assistant:\nI'm doing great. How can I help you today?### User:\nI'd like to show off how chat templating works!",
            },
            ChatTemplateTestItem {
                name: "Qwen/Qwen1.5-72B-Chat",
                chat_template: "{% for message in messages %}{% if loop.first and messages[0]['role'] != 'system' %}{{ '<|im_start|>system\\nYou are a helpful assistant<|im_end|>\\n' }}{% endif %}{{'<|im_start|>' + message['role'] + '\\n' + message['content']}}{% if (loop.last and add_generation_prompt) or not loop.last %}{{ '<|im_end|>' + '\\n'}}{% endif %}{% endfor %}{% if add_generation_prompt and messages[-1]['role'] != 'assistant' %}{{ '<|im_start|>assistant\\n' }}{% endif %}",
                input: ChatTemplateInputs {
                    messages: example_chat.clone(),
                    add_generation_prompt: false,
                    bos_token: Some("<s>"),
                    eos_token: Some("</s>"),
1475
                    ..Default::default()
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
                },
                target: "<|im_start|>system\nYou are a helpful assistant<|im_end|>\n<|im_start|>user\nHello, how are you?<|im_end|>\n<|im_start|>assistant\nI'm doing great. How can I help you today?<|im_end|>\n<|im_start|>user\nI'd like to show off how chat templating works!",
            },
            ChatTemplateTestItem {
                name: "deepseek-ai/deepseek-llm-7b-chat",
                chat_template: "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{{ bos_token }}{% for message in messages %}{% if message['role'] == 'user' %}{{ 'User: ' + message['content'] + '\\n\\n' }}{% elif message['role'] == 'assistant' %}{{ 'Assistant: ' + message['content'] + eos_token }}{% elif message['role'] == 'system' %}{{ message['content'] + '\\n\\n' }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ 'Assistant:' }}{% endif %}",
                input: ChatTemplateInputs {
                    messages: example_chat.clone(),
                    add_generation_prompt: false,
                    bos_token: Some("<|begin▁of▁sentence|>"),
                    eos_token: Some("<|end▁of▁sentence|>"),
1487
                    ..Default::default()
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
                },
                target: "<|begin▁of▁sentence|>User: Hello, how are you?\n\nAssistant: I'm doing great. How can I help you today?<|end▁of▁sentence|>User: I'd like to show off how chat templating works!\n\n",
            },
            ChatTemplateTestItem {
                name: "h2oai/h2o-danube-1.8b-chat",
                chat_template: "{% for message in messages %}{% if message['role'] == 'user' %}{{ '<|prompt|>' + message['content'] + eos_token }}{% elif message['role'] == 'system' %}{{ '<|system|>' + message['content'] + eos_token }}{% elif message['role'] == 'assistant' %}{{ '<|answer|>'  + message['content'] + eos_token }}{% endif %}{% if loop.last and add_generation_prompt %}{{ '<|answer|>' }}{% endif %}{% endfor %}",
                input: ChatTemplateInputs {
                    messages: example_chat.clone(),
                    add_generation_prompt: false,
                    bos_token: Some("<s>"),
                    eos_token: Some("</s>"),
1499
                    ..Default::default()
1500
                },
1501
                target: "<|prompt|>Hello, how are you?</s><|answer|>I'm doing great. How can I help you today?</s><|prompt|>I'd like to show off how chat templating works!</s>",
1502
1503
1504
1505
1506
1507
1508
1509
1510
            },
            ChatTemplateTestItem {
                name: "internlm/internlm2-chat-7b",
                chat_template: "{% if messages[0]['role'] == 'user' or messages[0]['role'] == 'system' %}{{ bos_token }}{% endif %}{% for message in messages %}{{ '<|im_start|>' + message['role'] + '\\n' + message['content'] + '<|im_end|>' + '\\n' }}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\\n' }}{% elif messages[-1]['role'] == 'assistant' %}{{ eos_token }}{% endif %}",
                input: ChatTemplateInputs {
                    messages: example_chat.clone(),
                    add_generation_prompt: false,
                    bos_token: Some("<s>"),
                    eos_token: Some("</s>"),
1511
                    ..Default::default()
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
                },
                target: "<s><|im_start|>user\nHello, how are you?<|im_end|>\n<|im_start|>assistant\nI'm doing great. How can I help you today?<|im_end|>\n<|im_start|>user\nI'd like to show off how chat templating works!<|im_end|>\n",
            },
            ChatTemplateTestItem {
                name: "TheBloke/deepseek-coder-33B-instruct-AWQ",
                chat_template: "{%- set found_item = false -%}\n{%- for message in messages -%}\n    {%- if message['role'] == 'system' -%}\n        {%- set found_item = true -%}\n    {%- endif -%}\n{%- endfor -%}\n{%- if not found_item -%}\n{{'You are an AI programming assistant, utilizing the Deepseek Coder model, developed by Deepseek Company, and you only answer questions related to computer science. For politically sensitive questions, security and privacy issues, and other non-computer science questions, you will refuse to answer.\\n'}}\n{%- endif %}\n{%- for message in messages %}\n    {%- if message['role'] == 'system' %}\n{{ message['content'] }}\n    {%- else %}\n        {%- if message['role'] == 'user' %}\n{{'### Instruction:\\n' + message['content'] + '\\n'}}\n        {%- else %}\n{{'### Response:\\n' + message['content'] + '\\n<|EOT|>\\n'}}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{{'### Response:\\n'}}\n",
                input: ChatTemplateInputs {
                    messages: example_chat.clone(),
                    add_generation_prompt: false,
                    bos_token: Some("<|begin▁of▁sentence|>"),
                    eos_token: Some("<|EOT|>"),
1523
                    ..Default::default()
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
                },
                target: "You are an AI programming assistant, utilizing the Deepseek Coder model, developed by Deepseek Company, and you only answer questions related to computer science. For politically sensitive questions, security and privacy issues, and other non-computer science questions, you will refuse to answer.\n### Instruction:\nHello, how are you?\n### Response:\nI'm doing great. How can I help you today?\n<|EOT|>\n### Instruction:\nI'd like to show off how chat templating works!\n### Response:\n",
            },
            ChatTemplateTestItem {
                name: "ericzzz/falcon-rw-1b-chat",
                // `.strip()` has been replaced with `| trim` in the following template
                chat_template: "{% for message in messages %}{% if loop.index > 1 and loop.previtem['role'] != 'assistant' %}{{ ' ' }}{% endif %}{% if message['role'] == 'system' %}{{ '[SYS] ' + message['content'] | trim }}{% elif message['role'] == 'user' %}{{ '[INST] ' + message['content'] | trim }}{% elif message['role'] == 'assistant' %}{{ '[RESP] '  + message['content'] + eos_token }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ ' [RESP] ' }}{% endif %}",
                input: ChatTemplateInputs {
                    messages: example_chat.clone(),
                    add_generation_prompt: false,
                    bos_token: Some("<|endoftext|>"),
                    eos_token: Some("<|endoftext|>"),
1536
                    ..Default::default()
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
                },
                target: "[INST] Hello, how are you? [RESP] I'm doing great. How can I help you today?<|endoftext|>[INST] I'd like to show off how chat templating works!",
            },
            ChatTemplateTestItem {
                name: "abacusai/Smaug-34B-v0.1",
                chat_template: "{%- for idx in range(0, messages|length) -%}\n{%- if messages[idx]['role'] == 'user' -%}\n{%- if idx > 1 -%}\n{{- bos_token + '[INST] ' + messages[idx]['content'] + ' [/INST]' -}}\n{%- else -%}\n{{- messages[idx]['content'] + ' [/INST]' -}}\n{%- endif -%}\n{% elif messages[idx]['role'] == 'system' %}\n{{- '[INST] <<SYS>>\\n' + messages[idx]['content'] + '\\n<</SYS>>\\n\\n' -}}\n{%- elif messages[idx]['role'] == 'assistant' -%}\n{{- ' '  + messages[idx]['content'] + ' ' + eos_token -}}\n{% endif %}\n{% endfor %}",
                input: ChatTemplateInputs {
                    messages: example_chat.clone(),
                    add_generation_prompt: false,
                    bos_token: Some("<s>"),
                    eos_token: Some("</s>"),
1548
                    ..Default::default()
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
                },
                target: "Hello, how are you? [/INST] I'm doing great. How can I help you today? </s><s>[INST] I'd like to show off how chat templating works! [/INST]",
            },
            ChatTemplateTestItem {
                name: "maywell/Synatra-Mixtral-8x7B",
                chat_template: "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n{% for message in messages %}{% if message['role'] == 'user' %}### Instruction:\n{{ message['content']|trim -}}{% if not loop.last %}{% endif %}\n{% elif message['role'] == 'assistant' %}### Response:\n{{ message['content']|trim -}}{% if not loop.last %}{% endif %}\n{% elif message['role'] == 'system' %}{{ message['content']|trim -}}{% if not loop.last %}{% endif %}\n{% endif %}\n{% endfor %}\n{% if add_generation_prompt and messages[-1]['role'] != 'assistant' %}\n### Response:\n{% endif %}",
                input: ChatTemplateInputs {
                    messages: example_chat.clone(),
                    add_generation_prompt: false,
                    bos_token: Some("<s>"),
                    eos_token: Some("</s>"),
1560
                    ..Default::default()
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
                },
                target: "Below is an instruction that describes a task. Write a response that appropriately completes the request.### Instruction:Hello, how are you?### Response:I'm doing great. How can I help you today?### Instruction:I'd like to show off how chat templating works!",
            },
            ChatTemplateTestItem {
                name: "deepseek-ai/deepseek-coder-33b-instruct",
                chat_template: "{% if not add_generation_prompt is defined %}\n{% set add_generation_prompt = false %}\n{% endif %}\n{%- set ns = namespace(found=false) -%}\n{%- for message in messages -%}\n    {%- if message['role'] == 'system' -%}\n        {%- set ns.found = true -%}\n    {%- endif -%}\n{%- endfor -%}\n{{bos_token}}{%- if not ns.found -%}\n{{'You are an AI programming assistant, utilizing the Deepseek Coder model, developed by Deepseek Company, and you only answer questions related to computer science. For politically sensitive questions, security and privacy issues, and other non-computer science questions, you will refuse to answer\\n'}}\n{%- endif %}\n{%- for message in messages %}\n    {%- if message['role'] == 'system' %}\n{{ message['content'] }}\n    {%- else %}\n        {%- if message['role'] == 'user' %}\n{{'### Instruction:\\n' + message['content'] + '\\n'}}\n        {%- else %}\n{{'### Response:\\n' + message['content'] + '\\n<|EOT|>\\n'}}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{% if add_generation_prompt %}\n{{'### Response:'}}\n{% endif %}",
                input: ChatTemplateInputs {
                    messages: example_chat.clone(),
                    add_generation_prompt: false,
                    bos_token: Some("<|begin▁of▁sentence|>"),
                    eos_token: Some("</EOT>"),
1572
                    ..Default::default()
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
                },
                target: "<|begin▁of▁sentence|>You are an AI programming assistant, utilizing the Deepseek Coder model, developed by Deepseek Company, and you only answer questions related to computer science. For politically sensitive questions, security and privacy issues, and other non-computer science questions, you will refuse to answer\n### Instruction:\nHello, how are you?\n### Response:\nI'm doing great. How can I help you today?\n<|EOT|>\n### Instruction:\nI'd like to show off how chat templating works!\n",
            },
            // NOT INCLUDED
            // - meetkai/functionary-medium-v2.2
            // - fireworks-ai/firefunction-v1
            // https://github
            ChatTemplateTestItem {
                name: "maywell/PiVoT-MoE",
                chat_template: "{{ (messages|selectattr('role', 'equalto', 'system')|list|last).content|trim if (messages|selectattr('role', 'equalto', 'system')|list) else '' }}{% for message in messages %}{% if message['role'] == 'system' %}{{ message['content']|trim }}{% elif message['role'] == 'user' %}### Instruction: {{ message['content']|trim }}{% elif message['role'] == 'assistant' %}### Response: {{ message['content']|trim }}{% elif message['role'] == 'user_context' %}### Input: {{ message['content']|trim }}{% endif %}{% if not loop.last %}\n{% endif %}{% endfor %}{% if add_generation_prompt and messages[-1]['role'] != 'assistant' %}### Response:{% endif %}",
                input: ChatTemplateInputs {
                    messages: example_chat_with_system.clone(),
                    add_generation_prompt: false,
                    bos_token: Some("<s>"),
                    eos_token: Some("</s>"),
1588
                    ..Default::default()
1589
1590
                },
                target: "You are a friendly chatbot who always responds in the style of a pirateYou are a friendly chatbot who always responds in the style of a pirate### Instruction: Hello, how are you?### Response: I'm doing great. How can I help you today?### Instruction: I'd like to show off how chat templating works!",
1591
            },
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
        ];

        #[allow(unused_variables)] // name is unused
        for ChatTemplateTestItem {
            name,
            chat_template,
            input,
            target,
        } in test_custom_templates
        {
            let mut env = Environment::new();
            env.add_function("raise_exception", raise_exception);
            // trim all the whitespace
            let chat_template = chat_template
                .lines()
                .map(|line| line.trim())
                .collect::<Vec<&str>>()
                .join("");

            let tmpl = env.template_from_str(&chat_template);
            let result = tmpl.unwrap().render(input).unwrap();
            assert_eq!(result, target);
        }
    }
1616
}