Update script.py

script.py

@@ -20,6 +20,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 from retrying import retry
 
 bnb_config = BitsAndBytesConfig(load_in_8bit=True)
+RunningInCOLAB = 'google.colab' in str(get_ipython()) if hasattr(__builtins__,'__IPYTHON__') else False
 
 # Load tokenizer
 _tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-7B-Instruct")
@@ -27,7 +28,7 @@ _tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-7B-Instruct")
 # Load model with quantization
 base_model = AutoModelForCausalLM.from_pretrained(
     "Qwen/Qwen2.5-Coder-7B-Instruct",
-    quantization_config=bnb_config,
+    quantization_config=bnb_config if RunningInCOLAB else None,
     device_map="auto",
 )