Skip to content

fix: make the wasm32 target compile again, and keep it compiling #156

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
lets-order-some-fries wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: make the wasm32 target compile again, and keep it compiling #156

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 36 additions & 0 deletions .github/workflows/wasm32.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
# Kept separate from CI.yml on purpose: that file is generated by
# `maturin generate-ci github`, so a job added there is lost on the next
# regeneration.
name: wasm32

on:
push:
branches:
- main
- master
pull_request:
workflow_dispatch:

permissions:
contents: read

jobs:
check:
name: cargo check (wasm32-unknown-unknown)
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4

- name: Setup Rust toolchain
uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable
with:
toolchain: stable
targets: wasm32-unknown-unknown

# The crate carries a cfg(target_arch = "wasm32") branch that nothing
# built, so it drifted out of compiling and stayed that way for a year.
# Checking the target is what stops that recurring.
- name: Check wasm32 target
shell: bash
run: |
cargo check --target wasm32-unknown-unknown --lib
118 changes: 55 additions & 63 deletions src/registry.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,42 +41,61 @@ impl std::fmt::Debug for HarmonyEncodingName {
}
}

/// Build the encoding around an already-loaded tokenizer.
///
/// The native and wasm32 loaders differ only in how they obtain the tokenizer
/// (blocking vs async), so everything else lives here. Keeping two copies of
/// this struct literal is what let the wasm32 branch drift: it went on setting
/// a `conversation_has_function_tools` field for a year after d00ac3d removed
/// it from `HarmonyEncoding`, and nothing built wasm32 to notice.
fn build_harmony_encoding(
name: HarmonyEncodingName,
tokenizer: crate::tiktoken::CoreBPE,
tokenizer_name: String,
) -> HarmonyEncoding {
let n_ctx = 1_048_576; // 2^20
let max_action_length = 524_288; // 2^19
HarmonyEncoding {
name: name.to_string(),
n_ctx,
tokenizer: Arc::new(tokenizer),
tokenizer_name,
max_message_tokens: n_ctx - max_action_length,
max_action_length,
format_token_mapping: make_mapping([
(FormattingToken::Start, "<|start|>"),
(FormattingToken::Message, "<|message|>"),
(FormattingToken::EndMessage, "<|end|>"),
(FormattingToken::EndMessageDoneSampling, "<|return|>"),
(FormattingToken::Refusal, "<|refusal|>"),
(FormattingToken::ConstrainedFormat, "<|constrain|>"),
(FormattingToken::Channel, "<|channel|>"),
(FormattingToken::EndMessageAssistantToTool, "<|call|>"),
(FormattingToken::BeginUntrusted, "<|untrusted|>"),
(FormattingToken::EndUntrusted, "<|end_untrusted|>"),
]),
stop_formatting_tokens: HashSet::from([
FormattingToken::EndMessageDoneSampling,
FormattingToken::EndMessageAssistantToTool,
FormattingToken::EndMessage,
]),
stop_formatting_tokens_for_assistant_actions: HashSet::from([
FormattingToken::EndMessageDoneSampling,
FormattingToken::EndMessageAssistantToTool,
]),
}
}

#[cfg(not(target_arch = "wasm32"))]
pub fn load_harmony_encoding(name: HarmonyEncodingName) -> anyhow::Result<HarmonyEncoding> {
match name {
HarmonyEncodingName::HarmonyGptOss => {
let n_ctx = 1_048_576; // 2^20
let max_action_length = 524_288; // 2^19
let encoding_ext = tiktoken_ext::Encoding::O200kHarmony;
Ok(HarmonyEncoding {
name: name.to_string(),
n_ctx,
tokenizer: Arc::new(encoding_ext.load()?),
tokenizer_name: encoding_ext.name().to_owned(),
max_message_tokens: n_ctx - max_action_length,
max_action_length,
format_token_mapping: make_mapping([
(FormattingToken::Start, "<|start|>"),
(FormattingToken::Message, "<|message|>"),
(FormattingToken::EndMessage, "<|end|>"),
(FormattingToken::EndMessageDoneSampling, "<|return|>"),
(FormattingToken::Refusal, "<|refusal|>"),
(FormattingToken::ConstrainedFormat, "<|constrain|>"),
(FormattingToken::Channel, "<|channel|>"),
(FormattingToken::EndMessageAssistantToTool, "<|call|>"),
(FormattingToken::BeginUntrusted, "<|untrusted|>"),
(FormattingToken::EndUntrusted, "<|end_untrusted|>"),
]),
stop_formatting_tokens: HashSet::from([
FormattingToken::EndMessageDoneSampling,
FormattingToken::EndMessageAssistantToTool,
FormattingToken::EndMessage,
]),
stop_formatting_tokens_for_assistant_actions: HashSet::from([
FormattingToken::EndMessageDoneSampling,
FormattingToken::EndMessageAssistantToTool,
]),
})
Ok(build_harmony_encoding(
name,
encoding_ext.load()?,
encoding_ext.name().to_owned(),
))
}
}
}
Expand All @@ -85,39 +104,12 @@ pub fn load_harmony_encoding(name: HarmonyEncodingName) -> anyhow::Result<Harmon
pub async fn load_harmony_encoding(name: HarmonyEncodingName) -> anyhow::Result<HarmonyEncoding> {
match name {
HarmonyEncodingName::HarmonyGptOss => {
let n_ctx = 1_048_576; // 2^20
let max_action_length = 524_288; // 2^19
let encoding_ext = tiktoken_ext::Encoding::O200kHarmony;
Ok(HarmonyEncoding {
name: name.to_string(),
n_ctx,
tokenizer: Arc::new(encoding_ext.load().await?),
tokenizer_name: encoding_ext.name().to_owned(),
max_message_tokens: n_ctx - max_action_length,
max_action_length,
format_token_mapping: make_mapping([
(FormattingToken::Start, "<|start|>"),
(FormattingToken::Message, "<|message|>"),
(FormattingToken::EndMessage, "<|end|>"),
(FormattingToken::EndMessageDoneSampling, "<|return|>"),
(FormattingToken::Refusal, "<|refusal|>"),
(FormattingToken::ConstrainedFormat, "<|constrain|>"),
(FormattingToken::Channel, "<|channel|>"),
(FormattingToken::EndMessageAssistantToTool, "<|call|>"),
(FormattingToken::BeginUntrusted, "<|untrusted|>"),
(FormattingToken::EndUntrusted, "<|end_untrusted|>"),
]),
stop_formatting_tokens: HashSet::from([
FormattingToken::EndMessageDoneSampling,
FormattingToken::EndMessageAssistantToTool,
FormattingToken::EndMessage,
]),
stop_formatting_tokens_for_assistant_actions: HashSet::from([
FormattingToken::EndMessageDoneSampling,
FormattingToken::EndMessageAssistantToTool,
]),
conversation_has_function_tools: Arc::new(AtomicBool::new(false)),
})
Ok(build_harmony_encoding(
name,
encoding_ext.load().await?,
encoding_ext.name().to_owned(),
))
}
}
}
Expand Down
50 changes: 44 additions & 6 deletions src/tiktoken_ext/public_encodings.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -390,19 +390,15 @@ where
load_tiktoken_vocab(reader, expected_hash)
}

pub fn load_encoding_from_file<P, S, TS>(
file_path: P,
expected_hash: Option<&str>,
fn build_core_bpe<S, TS>(
encoder: HashMap<Vec<u8>, Rank>,
special_tokens: S,
pattern: &str,
) -> Result<CoreBPE, LoadError>
where
P: AsRef<Path>,
S: IntoIterator<Item = (TS, Rank)>,
TS: Into<String>,
{
let encoder = load_tiktoken_vocab_file(file_path, expected_hash)
.map_err(LoadError::InvalidTiktokenVocabFile)?;
CoreBPE::new(
encoder,
special_tokens.into_iter().map(|(k, v)| (k.into(), v)),
Expand All @@ -411,6 +407,48 @@ where
.map_err(LoadError::CoreBPECreationFailed)
}

pub fn load_encoding_from_file<P, S, TS>(
file_path: P,
expected_hash: Option<&str>,
special_tokens: S,
pattern: &str,
) -> Result<CoreBPE, LoadError>
where
P: AsRef<Path>,
S: IntoIterator<Item = (TS, Rank)>,
TS: Into<String>,
{
let encoder = load_tiktoken_vocab_file(file_path, expected_hash)
.map_err(LoadError::InvalidTiktokenVocabFile)?;
build_core_bpe(encoder, special_tokens, pattern)
}

/// Same as [`load_encoding_from_file`], but reads the vocabulary from memory.
///
/// The wasm32 loader has always called this -- there is no filesystem to read
/// on that target -- but it was never defined, which is part of why wasm32 has
/// not compiled.
///
/// Gated to wasm32 because `public_encodings` is a private module, so an
/// unreachable `pub fn` trips `dead_code` under `-D warnings` on other targets.
/// Loading a vocabulary you already hold is useful off wasm32 too (an embedded
/// or air-gapped copy, cf. #101); exposing it there is a separate API decision.
#[cfg(target_arch = "wasm32")]
pub fn load_encoding_from_bytes<S, TS>(
vocab_bytes: &[u8],
expected_hash: Option<&str>,
special_tokens: S,
pattern: &str,
) -> Result<CoreBPE, LoadError>
where
S: IntoIterator<Item = (TS, Rank)>,
TS: Into<String>,
{
let encoder = load_tiktoken_vocab(vocab_bytes, expected_hash)
.map_err(LoadError::InvalidTiktokenVocabFile)?;
build_core_bpe(encoder, special_tokens, pattern)
}

/// This returns the path to a file containing the data at `url`. If the file is
/// cached, it is used. Otherwise, the file is downloaded and cached.
#[cfg(not(target_arch = "wasm32"))]
Expand Down