feat(templates): avanterules filetype support (closes #254) (#466)

Signed-off-by: Aaron Pham <contact@aarnphm.xyz>
This commit is contained in:
Aaron Pham
2024-09-03 04:09:13 -04:00
committed by GitHub
parent 054695cc63
commit 4ad913435c
31 changed files with 962 additions and 265 deletions

View File

@@ -68,13 +68,12 @@ fn encode(state: &State, text: String) -> LuaResult<(Vec<usize>, usize, usize)>
}
}
fn from_pretrained(state: &State, model: String) -> LuaResult<()> {
fn from_pretrained(state: &State, model: String) {
let mut tokenizer_mutex = state.tokenizer.lock().unwrap();
*tokenizer_mutex = Some(match model.as_str() {
"gpt-4o" => TokenizerType::Tiktoken(Tiktoken::new(model)),
_ => TokenizerType::HuggingFace(HuggingFaceTokenizer::new(model)),
});
Ok(())
}
#[mlua::lua_module]
@@ -86,7 +85,10 @@ fn avante_tokenizers(lua: &Lua) -> LuaResult<LuaTable> {
let exports = lua.create_table()?;
exports.set(
"from_pretrained",
lua.create_function(move |_, model: String| from_pretrained(&state, model))?,
lua.create_function(move |_, model: String| {
from_pretrained(&state, model);
Ok(())
})?,
)?;
exports.set(
"encode",