huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-11 20:34:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

Using @drbh patch.

Browse Source
This commit is contained in:
Nicolas Patry 2024-04-23 14:27:43 +00:00
parent af24703708
commit 429092683b
1 changed files with 12 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

35
router/src/main.rs
View File

@ -187,12 +187,12 @@ async fn main() -> Result<(), RouterError> {
None, None,
} }
let api = if use_api { let api = if use_api {
tracing::info!("Using the Hugging Face API");
if std::env::var("HF_HUB_OFFLINE") == Ok("1".to_string()) { if std::env::var("HF_HUB_OFFLINE") == Ok("1".to_string()) {
let cache = Cache::default(); let cache = Cache::default();
tracing::warn!("Offline mode active using cache defaults"); tracing::warn!("Offline mode active using cache defaults");
Type::Cache(cache) Type::Cache(cache)
} else { } else {
tracing::info!("Using the Hugging Face API");
match api_builder().build() { match api_builder().build() {
Ok(api) => Type::Api(api), Ok(api) => Type::Api(api),
Err(_) => { Err(_) => {
@ -207,18 +207,12 @@ async fn main() -> Result<(), RouterError> {
// Load tokenizer and model info // Load tokenizer and model info
let (tokenizer_filename, config_filename, tokenizer_config_filename, model_info) = match api { let (tokenizer_filename, config_filename, tokenizer_config_filename, model_info) = match api {
Type::None => { Type::None => (
let tokenizer_filename = Some(local_path.join("tokenizer.json")); Some(local_path.join("tokenizer.json")),
let config_filename = Some(local_path.join("config.json")); Some(local_path.join("config.json")),
let tokenizer_config_filename = Some(local_path.join("tokenizer_config.json")); Some(local_path.join("tokenizer_config.json")),
let model_info = None; None,
( ),
tokenizer_filename,
config_filename,
tokenizer_config_filename,
model_info,
)
}
Type::Api(api) => { Type::Api(api) => {
let api_repo = api.repo(Repo::with_revision( let api_repo = api.repo(Repo::with_revision(
tokenizer_name.to_string(), tokenizer_name.to_string(),
@ -247,21 +241,16 @@ async fn main() -> Result<(), RouterError> {
) )
} }
Type::Cache(cache) => { Type::Cache(cache) => {
let cache_repo = cache.repo(Repo::with_revision( let repo = cache.repo(Repo::with_revision(
tokenizer_name.to_string(), tokenizer_name.to_string(),
RepoType::Model, RepoType::Model,
revision.clone().unwrap_or_else(|| "main".to_string()), revision.clone().unwrap_or_else(|| "main".to_string()),
)); ));
let tokenizer_filename = cache_repo.get("tokenizer.json");
let config_filename = cache_repo.get("config.json");
let tokenizer_config_filename = cache_repo.get("tokenizer_config.json");
let model_info = None;
( (
tokenizer_filename, repo.get("tokenizer.json"),
config_filename, repo.get("config.json"),
tokenizer_config_filename, repo.get("tokenizer_config.json"),
model_info, None,
) )
} }
}; };