rust 获取 hugging face 模型 chat template

173 阅读1分钟

前言

rust 版 tokenizer 没有直接获取 chat template的方法,记录一下探索获取它的方法

方案

手动复制保存

直接在模型库右边点Chat template -> Copy,保存下来 image.png

image.png

爬虫

通过浏览器查看Chat template元素时发现div.SVELTE_HYDRATER.contents存了chat template相关信息,于是可以用 reqwestscraper 解析 html 取 div.SVELTE_HYDRATER.contentsdata-props 属性转成 json 再过滤取到 chat template

image.png

let repo = "Qwen/Qwen2.5-7B-Instruct";

let resp = Client::new()
    .get(format!("https://huggingface.co/{repo}"))
    .send()
    .await?;

let document = Html::parse_document(&resp.text().await?);
let selector = Selector::parse("div.SVELTE_HYDRATER.contents").unwrap();

for elem in document.select(&selector) {
    if let Some(attr) = elem.attr("data-props") {
        let attr: Value = serde_json::from_str(attr)?;
        println!("{:#?}", attr);
    }
}

api

hugging face api没有提供chat template,但是可以通过api获取模型的配置信息再得到template

let repo = "Qwen/Qwen2.5-7B-Instruct";
let repo = Api::new()?.model(repo.into());

let json: Value = repo.info_request().send().await?.json().await?;

let chat_template = json["config"]["tokenizer_config"]["chat_template"].as_str().unwrap();

终极方案

最后发现原来模型库的tokenizer_config.json就有chat template信息,完全可以通过api下载这个文件获取

image.png

async fn load_template(tokenizer_repo: &str) -> Result<Value> {
    let pth = Api::new()?
        .model(tokenizer_repo.to_string())
        .get("tokenizer_config.json")
        .await?;

    let file = File::open(pth)?;
    let mut json: Value = serde_json::from_reader(BufReader::new(file))?;

    Ok(json["chat_template"].take())
}