Skip to content

Commit af52be1

Browse files
authored
fix context size for fine-tuned models (#109)
1 parent 7c20dc6 commit af52be1

File tree

2 files changed

+17
-0
lines changed

2 files changed

+17
-0
lines changed

tiktoken-rs/src/model.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,10 @@ macro_rules! starts_with_any {
3232
///
3333
/// This function does not panic. It returns a default value of 4096 if the model is not recognized.
3434
pub fn get_context_size(model: &str) -> usize {
35+
if let Some(rest) = model.strip_prefix("ft:") {
36+
let base = rest.split(':').next().unwrap_or(rest);
37+
return get_context_size(base);
38+
}
3539
if starts_with_any!(model, "o1-") {
3640
return 128_000;
3741
}

tiktoken-rs/tests/model.rs

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
use tiktoken_rs::model::get_context_size;
2+
3+
#[test]
4+
fn test_finetuned_context_size() {
5+
assert_eq!(
6+
get_context_size("ft:gpt-3.5-turbo-0125:custom"),
7+
get_context_size("gpt-3.5-turbo-0125")
8+
);
9+
assert_eq!(
10+
get_context_size("ft:gpt-4o:custom"),
11+
get_context_size("gpt-4o")
12+
);
13+
}

0 commit comments

Comments
 (0)