Upload app.py

app.py

@@ -140,9 +140,11 @@ class Social_Media_Captioner:
     def _load_model(self):
         self.bnb_config = BitsAndBytesConfig(
             load_in_4bit = True,
+            llm_int8_enable_fp32_cpu_offload=True,
             bnb_4bit_use_double_quant = True,
             bnb_4bit_quant_type= "nf4",
             bnb_4bit_compute_dtype=torch.bfloat16,
+            load_in_8bit_fp32_cpu_offload=True
             )
         self.model = AutoModelForCausalLM.from_pretrained(
             self.MODEL_NAME,
@@ -155,32 +157,32 @@ class Social_Media_Captioner:
         self.tokenizer = AutoTokenizer.from_pretrained(self.MODEL_NAME)
         self.tokenizer.pad_token = self.tokenizer.eos_token
 
-        if self.use_finetuned:
-            # LORA Config Model
-            self.lora_config = LoraConfig(
-                r=16,
-                lora_alpha=32,
-                target_modules=["query_key_value"],
-                lora_dropout=0.05,
-                bias="none",
-                task_type="CAUSAL_LM"
-            )
-            self.model = get_peft_model(self.model, self.lora_config)
-
-            # Fitting the adapters
-            self.peft_config = PeftConfig.from_pretrained(self.peft_model_name)
-            self.model = AutoModelForCausalLM.from_pretrained(
-                self.peft_config.base_model_name_or_path,
-                return_dict = True,
-                quantization_config = self.bnb_config,
-                device_map= "auto",
-                trust_remote_code = True
-                )
-            self.model = PeftModel.from_pretrained(self.model, self.peft_model_name)
-
-            # Defining the tokenizers
-            self.tokenizer = AutoTokenizer.from_pretrained(self.peft_config.base_model_name_or_path)
-            self.tokenizer.pad_token = self.tokenizer.eos_token
+        # if self.use_finetuned:
+        #     # LORA Config Model
+        #     self.lora_config = LoraConfig(
+        #         r=16,
+        #         lora_alpha=32,
+        #         target_modules=["query_key_value"],
+        #         lora_dropout=0.05,
+        #         bias="none",
+        #         task_type="CAUSAL_LM"
+        #     )
+        #     self.model = get_peft_model(self.model, self.lora_config)
+
+        #     # Fitting the adapters
+        #     self.peft_config = PeftConfig.from_pretrained(self.peft_model_name)
+        #     self.model = AutoModelForCausalLM.from_pretrained(
+        #         self.peft_config.base_model_name_or_path,
+        #         return_dict = True,
+        #         quantization_config = self.bnb_config,
+        #         device_map= "auto",
+        #         trust_remote_code = True
+        #         )
+        #     self.model = PeftModel.from_pretrained(self.model, self.peft_model_name)
+
+        #     # Defining the tokenizers
+        #     self.tokenizer = AutoTokenizer.from_pretrained(self.peft_config.base_model_name_or_path)
+        #     self.tokenizer.pad_token = self.tokenizer.eos_token
 
         self.model_loaded = True
         print("Model Loaded successfully")