dwb2023 commited on
Commit
f4720e3
·
verified ·
1 Parent(s): b1703c6

Update app.py

Browse files

adjust quantization_config

Files changed (1) hide show
  1. app.py +7 -2
app.py CHANGED
@@ -23,8 +23,13 @@ YT_LENGTH_LIMIT_S = 4800 # 1 hour 20 minutes
23
  device = 0 if torch.cuda.is_available() else "cpu"
24
 
25
  # Load model with bitsandbytes quantization
26
- bnb_config = bnb.QuantizationConfig(bits=4)
27
- pipe = pipeline(task="automatic-speech-recognition", model=MODEL_NAME, chunk_length_s=30, device=device, quantization_config=bnb_config)
 
 
 
 
 
28
 
29
  # Define paths and create directory if not exists
30
  JSON_DATASET_DIR = Path("json_dataset")
 
23
  device = 0 if torch.cuda.is_available() else "cpu"
24
 
25
  # Load model with bitsandbytes quantization
26
+ quantization_config = BitsAndBytesConfig(
27
+ load_in_4bit=True,
28
+ bnb_4bit_compute_dtype=torch.float16,
29
+ )
30
+
31
+ # bnb_config = bnb.QuantizationConfig(bits=4)
32
+ pipe = pipeline(task="automatic-speech-recognition", model=MODEL_NAME, chunk_length_s=30, device=device, quantization_config=quantization_config)
33
 
34
  # Define paths and create directory if not exists
35
  JSON_DATASET_DIR = Path("json_dataset")