lib.rs 8.12 KB
Newer Older
1
mod health;
2
/// Text Generation Inference Webserver
3
mod infer;
4
mod queue;
Olivier Dehaene's avatar
Olivier Dehaene committed
5
pub mod server;
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
6
mod validation;
Olivier Dehaene's avatar
Olivier Dehaene committed
7

8
use infer::Infer;
9
use queue::{Entry, Queue};
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
10
use serde::{Deserialize, Serialize};
11
use utoipa::ToSchema;
Olivier Dehaene's avatar
Olivier Dehaene committed
12
use validation::Validation;
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
13

14
15
/// Hub type
#[derive(Clone, Debug, Deserialize)]
16
pub struct HubModelInfo {
17
18
19
20
21
22
23
24
    #[serde(rename(deserialize = "id"))]
    pub model_id: String,
    pub sha: Option<String>,
    pub pipeline_tag: Option<String>,
}

#[derive(Clone, Debug, Serialize, ToSchema)]
pub struct Info {
25
    /// Model info
26
27
28
29
    #[schema(example = "bigscience/blomm-560m")]
    pub model_id: String,
    #[schema(nullable = true, example = "e985a63cdc139290c5f700ff1929f0b5942cced2")]
    pub model_sha: Option<String>,
30
31
32
33
    #[schema(example = "torch.float16")]
    pub model_dtype: String,
    #[schema(example = "cuda")]
    pub model_device_type: String,
34
35
    #[schema(nullable = true, example = "text-generation")]
    pub model_pipeline_tag: Option<String>,
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
    /// 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")]
    pub max_input_length: usize,
    #[schema(example = "2048")]
    pub max_total_tokens: usize,
    #[schema(example = "1.2")]
    pub waiting_served_ratio: f32,
    #[schema(example = "32000")]
    pub max_batch_total_tokens: u32,
    #[schema(example = "20")]
    pub max_waiting_tokens: usize,
    #[schema(example = "2")]
    pub validation_workers: usize,
    /// Router Info
56
57
58
59
60
61
    #[schema(example = "0.5.0")]
    pub version: &'static str,
    #[schema(nullable = true, example = "null")]
    pub sha: Option<&'static str>,
}

62
#[derive(Clone, Debug, Deserialize, ToSchema)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
63
pub(crate) struct GenerateParameters {
64
65
66
    #[serde(default)]
    #[schema(exclusive_minimum = 0, nullable = true, default = "null", example = 1)]
    pub best_of: Option<usize>,
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
    #[serde(default)]
    #[schema(
        exclusive_minimum = 0.0,
        nullable = true,
        default = "null",
        example = 0.5
    )]
    pub temperature: Option<f32>,
    #[serde(default)]
    #[schema(
        exclusive_minimum = 0.0,
        nullable = true,
        default = "null",
        example = 1.03
    )]
    pub repetition_penalty: Option<f32>,
    #[serde(default)]
    #[schema(exclusive_minimum = 0, nullable = true, default = "null", example = 10)]
    pub top_k: Option<i32>,
    #[serde(default)]
    #[schema(
        exclusive_minimum = 0.0,
        maximum = 1.0,
        nullable = true,
        default = "null",
        example = 0.95
    )]
    pub top_p: Option<f32>,
95
    #[serde(default)]
96
97
98
99
100
101
102
103
104
    #[schema(
        exclusive_minimum = 0.0,
        maximum = 1.0,
        nullable = true,
        default = "null",
        example = 0.95
    )]
    pub typical_p: Option<f32>,
    #[serde(default)]
105
    #[schema(default = "false", example = true)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
106
107
    pub do_sample: bool,
    #[serde(default = "default_max_new_tokens")]
108
    #[schema(exclusive_minimum = 0, exclusive_maximum = 512, default = "20")]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
109
    pub max_new_tokens: u32,
OlivierDehaene's avatar
OlivierDehaene committed
110
    #[serde(default)]
111
    #[schema(nullable = true, default = "null", example = false)]
112
113
    pub return_full_text: Option<bool>,
    #[serde(default)]
114
    #[schema(inline, max_items = 4, example = json ! (["photographer"]))]
115
    pub stop: Vec<String>,
OlivierDehaene's avatar
OlivierDehaene committed
116
    #[serde(default)]
117
    #[schema(nullable = true, default = "null", example = "null")]
118
119
    pub truncate: Option<usize>,
    #[serde(default)]
120
121
122
    #[schema(default = "false", example = true)]
    pub watermark: bool,
    #[serde(default)]
123
    #[schema(default = "true")]
OlivierDehaene's avatar
OlivierDehaene committed
124
    pub details: bool,
125
    #[serde(default)]
126
127
128
129
130
131
    #[schema(
        exclusive_minimum = 0,
        nullable = true,
        default = "null",
        example = "null"
    )]
132
    pub seed: Option<u64>,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
133
134
135
136
137
138
139
140
}

fn default_max_new_tokens() -> u32 {
    20
}

fn default_parameters() -> GenerateParameters {
    GenerateParameters {
141
        best_of: None,
142
143
144
145
        temperature: None,
        repetition_penalty: None,
        top_k: None,
        top_p: None,
146
        typical_p: None,
147
        do_sample: false,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
148
        max_new_tokens: default_max_new_tokens(),
149
        return_full_text: None,
150
        stop: Vec::new(),
151
        truncate: None,
152
        watermark: false,
OlivierDehaene's avatar
OlivierDehaene committed
153
        details: false,
154
        seed: None,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
155
156
157
    }
}

158
#[derive(Clone, Debug, Deserialize, ToSchema)]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
159
pub(crate) struct GenerateRequest {
160
    #[schema(example = "My name is Olivier and I")]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
161
162
163
164
165
    pub inputs: String,
    #[serde(default = "default_parameters")]
    pub parameters: GenerateParameters,
}

166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
#[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)]
    #[allow(dead_code)]
    pub stream: bool,
}

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

186
187
188
189
190
191
#[derive(Debug, Serialize, ToSchema)]
pub struct PrefillToken {
    #[schema(example = 0)]
    id: u32,
    #[schema(example = "test")]
    text: String,
192
    #[schema(nullable = true, example = - 0.34)]
193
194
195
    logprob: f32,
}

196
197
198
199
200
201
#[derive(Debug, Serialize, ToSchema)]
pub struct Token {
    #[schema(example = 0)]
    id: u32,
    #[schema(example = "test")]
    text: String,
202
    #[schema(nullable = true, example = - 0.34)]
203
    logprob: f32,
204
205
    #[schema(example = "false")]
    special: bool,
206
207
208
209
210
211
212
213
214
215
216
217
218
}

#[derive(Serialize, ToSchema)]
#[serde(rename_all(serialize = "snake_case"))]
pub(crate) enum FinishReason {
    #[schema(rename = "length")]
    Length,
    #[serde(rename = "eos_token")]
    #[schema(rename = "eos_token")]
    EndOfSequenceToken,
    #[schema(rename = "stop_sequence")]
    StopSequence,
}
219

220
221
222
223
224
225
226
227
228
229
230
231
232
233
#[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>,
}

234
#[derive(Serialize, ToSchema)]
OlivierDehaene's avatar
OlivierDehaene committed
235
pub(crate) struct Details {
236
237
238
    #[schema(example = "length")]
    pub finish_reason: FinishReason,
    #[schema(example = 1)]
OlivierDehaene's avatar
OlivierDehaene committed
239
    pub generated_tokens: u32,
240
    #[schema(nullable = true, example = 42)]
241
    pub seed: Option<u64>,
242
243
    pub prefill: Vec<PrefillToken>,
    pub tokens: Vec<Token>,
244
245
    #[serde(skip_serializing_if = "Option::is_none")]
    pub best_of_sequences: Option<Vec<BestOfSequence>>,
OlivierDehaene's avatar
OlivierDehaene committed
246
247
}

248
#[derive(Serialize, ToSchema)]
249
pub(crate) struct GenerateResponse {
250
    #[schema(example = "test")]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
251
    pub generated_text: String,
OlivierDehaene's avatar
OlivierDehaene committed
252
253
    #[serde(skip_serializing_if = "Option::is_none")]
    pub details: Option<Details>,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
254
}
255

256
257
258
259
260
261
#[derive(Serialize, ToSchema)]
pub(crate) struct StreamDetails {
    #[schema(example = "length")]
    pub finish_reason: FinishReason,
    #[schema(example = 1)]
    pub generated_tokens: u32,
262
    #[schema(nullable = true, example = 42)]
263
264
265
266
    pub seed: Option<u64>,
}

#[derive(Serialize, ToSchema)]
267
268
pub(crate) struct StreamResponse {
    pub token: Token,
269
    #[schema(nullable = true, default = "null", example = "test")]
270
    pub generated_text: Option<String>,
271
272
    #[schema(nullable = true, default = "null")]
    pub details: Option<StreamDetails>,
273
274
}

275
#[derive(Serialize, ToSchema)]
276
277
pub(crate) struct ErrorResponse {
    pub error: String,
278
    pub error_type: String,
279
}
280
281

#[cfg(test)]
282
mod tests {
283
284
285
    use std::io::Write;
    use tokenizers::Tokenizer;

286
287
288
289
290
291
292
293
294
    pub(crate) async fn get_tokenizer() -> Tokenizer {
        if !std::path::Path::new("tokenizer.json").exists() {
            let content = reqwest::get("https://huggingface.co/gpt2/raw/main/tokenizer.json")
                .await
                .unwrap()
                .bytes()
                .await
                .unwrap();
            let mut file = std::fs::File::create("tokenizer.json").unwrap();
295
296
297
298
299
            file.write_all(&content).unwrap();
        }
        Tokenizer::from_file("tokenizer.json").unwrap()
    }
}