mkzero commited on
Commit
def067f
·
verified ·
1 Parent(s): 147775b

feat: enhance quickstart_demo for Apple Silicon and auto hub load

Browse files
Files changed (1) hide show
  1. quickstart_demo.py +6 -4
quickstart_demo.py CHANGED
@@ -2,14 +2,16 @@
2
  Cerebellum-2B: 25ms Non-Autoregressive Agent Decision Model
3
  Quickstart Demo (3 lines to make a fast decision)
4
  """
 
 
5
  import torch
6
  from modeling_cerebellum import CerebellumModel
7
 
8
  def main():
9
- # 1. Load Model (BF16 or FP8)
10
- device = "cuda:0" if torch.cuda.is_available() else "cpu"
11
- print(f"Loading Cerebellum-2B on {device}...")
12
- model = CerebellumModel.from_pretrained(".", device=device)
13
 
14
  # 2. Define Agent State and Candidate Actions
15
  state = """
 
2
  Cerebellum-2B: 25ms Non-Autoregressive Agent Decision Model
3
  Quickstart Demo (3 lines to make a fast decision)
4
  """
5
+ import os
6
+ import sys
7
  import torch
8
  from modeling_cerebellum import CerebellumModel
9
 
10
  def main():
11
+ # 1. Load Model (Support local folder or Hugging Face Hub, auto-detect CUDA/MPS/CPU)
12
+ model_path = sys.argv[1] if len(sys.argv) > 1 else ("." if os.path.exists("./model.safetensors") else "mkzero/Cerebellum-2B-INT8")
13
+ print(f"Loading Cerebellum-2B from: {model_path}...")
14
+ model = CerebellumModel.from_pretrained(model_path)
15
 
16
  # 2. Define Agent State and Candidate Actions
17
  state = """