Update app.py

app.py

@@ -21,21 +21,14 @@ YT_LENGTH_LIMIT_S = 4800  # 1 hour 20 minutes
 
 device = 0 if torch.cuda.is_available() else "cpu"
 
-
-# Quantization
-bnb_config=BitsAndBytesConfig(
-    load_in_4bit=True
-)
+# Load model with bitsandbytes quantization
+bnb_config = BitsAndBytesConfig(load_in_4bit=True)
 
 # Load the model
-model = WhisperForConditionalGeneration.from_pretrained(
-    MODEL_NAME, load_in_4bit=True, device_map="auto"
-)
-
-tokenizer = WhisperTokenizer.from_pretrained(MODEL_NAME)
+model = model_class.from_pretrained(model_name, config=bnb_config)
 
-# Initialize the pipeline with the quantized model
-pipe = pipeline(task="automatic-speech-recognition", model=model, tokenizer=tokenizer, chunk_length_s=30, device=device)
+# bnb_config = bnb.QuantizationConfig(bits=4)
+pipe = pipeline(task="automatic-speech-recognition", model=model, chunk_length_s=30, device=device)
 
 # Define paths and create directory if not exists
 JSON_DATASET_DIR = Path("json_dataset")