generate.proto 4.2 KB
Newer Older
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
1
2
3
4
syntax = "proto3";

package generate.v1;

Olivier Dehaene's avatar
Olivier Dehaene committed
5
service TextGenerationService {
6
7
    /// Model Info
    rpc Info (InfoRequest) returns (InfoResponse) {}
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
8
    /// Service discovery
Olivier Dehaene's avatar
Olivier Dehaene committed
9
    rpc ServiceDiscovery (ServiceDiscoveryRequest) returns (ServiceDiscoveryResponse) {}
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
10
    /// Empties batch cache
Olivier Dehaene's avatar
Olivier Dehaene committed
11
    rpc ClearCache (ClearCacheRequest) returns (ClearCacheResponse);
12
13
    /// Remove requests from a cached batch
    rpc FilterBatch (FilterBatchRequest) returns (FilterBatchResponse);
14
15
16
17
    /// Prefill batch and decode first token
    rpc Prefill (PrefillRequest) returns (PrefillResponse);
    /// Decode token for a list of prefilled batches
    rpc Decode (DecodeRequest) returns (DecodeResponse);
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
18
19
}

20
21
22
23
24
25
26
27
28
/// Empty request
message InfoRequest {}

message InfoResponse {
    bool requires_padding = 1;
    string dtype = 2;
    string device_type = 3;
}

Olivier Dehaene's avatar
Olivier Dehaene committed
29
30
31
/// Empty request
message ServiceDiscoveryRequest {}

Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
32
message ServiceDiscoveryResponse {
Olivier Dehaene's avatar
Olivier Dehaene committed
33
    /// Other shards urls
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
34
35
36
    repeated string urls = 1;
}

37
38
39
40
message ClearCacheRequest {
    /// Optional batch id
    optional uint64 id = 1;
}
Olivier Dehaene's avatar
Olivier Dehaene committed
41
42
43
44

/// Empty response
message ClearCacheResponse {}

OlivierDehaene's avatar
OlivierDehaene committed
45
message NextTokenChooserParameters {
46
    /// exponential scaling output probability distribution
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
47
    float temperature = 1;
48
    /// restricting to the k highest probability elements
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
49
    uint32 top_k = 2;
50
    /// restricting to top tokens summing to prob_cut_off <= prob_cut_off
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
51
    float top_p = 3;
52
53
    /// restricting to top tokens summing to prob_cut_off <= prob_cut_off
    float typical_p = 4;
54
    /// apply sampling on the logits
55
    bool do_sample = 5;
56
    /// random seed for sampling
57
    uint64 seed = 6;
58
    /// repetition penalty
59
    float repetition_penalty = 7;
60
    /// token watermarking using "A Watermark for Large Language Models"
61
    bool watermark = 8;
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
62
63
}

64
65
66
67
68
message StoppingCriteriaParameters {
    /// Maximum number of generated tokens
    uint32 max_new_tokens = 1;
    /// Optional stopping sequences
    repeated string stop_sequences = 2;
69
70
71
    /// Ignore end of sequence token
    /// used for benchmarking
    bool ignore_eos_token = 3;
72
73
}

Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
74
75
76
77
78
message Request {
    /// Request ID
    uint64 id = 1;
    /// The generation context
    string inputs = 2;
79
80
    /// Context truncation
    uint32 truncate = 3;
OlivierDehaene's avatar
OlivierDehaene committed
81
    /// Next Token Chooser Parameters
82
    NextTokenChooserParameters parameters = 4;
83
    /// Stopping Criteria Parameters
84
    StoppingCriteriaParameters stopping_parameters = 5;
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
85
86
87
88
89
90
91
}

message Batch {
    /// Batch ID
    uint64 id = 1;
    /// Individual requests
    repeated Request requests = 2;
Olivier Dehaene's avatar
Olivier Dehaene committed
92
93
    /// Batch size (==len(requests))
    uint32 size = 3;
94
95
    /// Maximum number of tokens this batch will grow to
    uint32 max_tokens = 4;
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
96
97
}

98
99
100
101
102
103
enum FinishReason {
    FINISH_REASON_LENGTH = 0;
    FINISH_REASON_EOS_TOKEN = 1;
    FINISH_REASON_STOP_SEQUENCE = 2;
}

Olivier Dehaene's avatar
Olivier Dehaene committed
104
message GeneratedText {
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
105
    /// Output
106
    string text = 1;
107
    /// Number of generated tokens
108
    uint32 generated_tokens = 2;
109
    /// Finish reason
110
    FinishReason finish_reason = 3;
111
    /// Seed
112
    optional uint64 seed = 4;
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
113
114
}

115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
message PrefillTokens {
    /// Prefill Token IDs
    repeated uint32 ids = 1;
    /// Prefill Logprobs
    repeated float logprobs = 2;
    /// Prefill tokens
    repeated string texts = 3;
}

message Generation {
    /// Request ID
    uint64 request_id = 1;
    /// Prefill tokens (optional)
    PrefillTokens prefill_tokens = 2;
    /// Token ID
    uint32 token_id = 3;
    /// Logprob
    float token_logprob = 4;
    /// Text
    string token_text = 5;
135
136
    /// Is it a special token
    bool token_is_special = 6;
137
    /// Complete generated text
138
    GeneratedText generated_text = 7;
139
140
}

141
142
143
144
145
146
147
148
149
150
151
152
153
message FilterBatchRequest {
    /// Batch ID
    uint64 batch_id = 1;
    /// Requests to keep
    repeated Request keep_requests = 2;
}

message FilterBatchResponse {
    /// Filtered Batch (cached)
    Batch batch = 1;
}


154
message PrefillRequest {
Olivier Dehaene's avatar
Olivier Dehaene committed
155
156
    /// Batch
    Batch batch = 1;
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
157
158
}

159
160
161
message PrefillResponse {
    /// Generation
    repeated Generation generations = 1;
Olivier Dehaene's avatar
Olivier Dehaene committed
162
163
    /// Next batch (cached)
    optional Batch batch = 2;
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
164
165
}

166
message DecodeRequest {
Olivier Dehaene's avatar
Olivier Dehaene committed
167
168
169
    /// Cached batches
    repeated Batch batches = 1;
}
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
170

171
172
173
message DecodeResponse {
    /// Decodes
    repeated Generation generations = 1;
Olivier Dehaene's avatar
Olivier Dehaene committed
174
175
    /// Next batch (cached)
    optional Batch batch = 2;
176
}