Use resharded StableBeluga2

petals-infra · Aug 24, 2023 · 08f21f5 · 08f21f5
1 parent 3558732
commit 08f21f5
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/app.py b/app.py
@@ -27,7 +27,9 @@
  )
  model = model.to(config.DEVICE)
 
- model_name = model_info.adapter if model_info.adapter is not None else model_info.repo
+ model_name = model_info.name
+ if model_name is None: # Use default name based on model/repo repo
+ model_name = model_info.adapter if model_info.adapter is not None else model_info.repo
  models[model_name] = model, tokenizer
 
 logger.info("Starting Flask app")

diff --git a/config.py b/config.py
@@ -11,10 +11,11 @@
 class ModelInfo:
  repo: str
  adapter: Optional[str] = None
+ name: Optional[str] = None
 
 
 MODELS = [
- ModelInfo(repo="stabilityai/StableBeluga2"),
+ ModelInfo(repo="petals-team/StableBeluga2", name="stabilityai/StableBeluga2"),
  ModelInfo(repo="meta-llama/Llama-2-70b-chat-hf"),
  ModelInfo(repo="huggyllama/llama-65b", adapter="timdettmers/guanaco-65b"),
  ModelInfo(repo="huggyllama/llama-65b"),