Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 from peft import PeftModel
-from transformers import RobertaTokenizer, T5ForConditionalGeneration, AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from transformers import RobertaTokenizer, T5ForConditionalGeneration, AutoTokenizer, AutoModelForCausalLM
 import torch
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -13,11 +13,11 @@ model_id = {
 }
 
 # Quantization Config
-quantization_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.bfloat16
-)
+# quantization_config = BitsAndBytesConfig(
+#     load_in_4bit=True,
+#     bnb_4bit_quant_type="nf4",
+#     bnb_4bit_compute_dtype=torch.bfloat16
+# )
 
 # Load CodeT5 models
 codeT5_tokenizer = RobertaTokenizer.from_pretrained(model_id['CodeT5'])
@@ -30,14 +30,14 @@ fine_tuned_codeT5.eval()
 # Load StarCoder
 StarCoder_tokenizer = AutoTokenizer.from_pretrained(model_id['StarCoder'])
 StarCoder = AutoModelForCausalLM.from_pretrained(
-    model_id['StarCoder'], quantization_config=quantization_config, device_map="auto"
+    model_id['StarCoder'], device_map="auto"
 )
 StarCoder.eval()
 
 # Load CodeLlama
 CodeLlama_tokenizer = AutoTokenizer.from_pretrained(model_id['CodeLlama'])
 CodeLlama = AutoModelForCausalLM.from_pretrained(
-    model_id['CodeLlama'], quantization_config=quantization_config, device_map="auto"
+    model_id['CodeLlama'], device_map="auto"
 )
 CodeLlama.eval()