Upload app.py

app.py

@@ -156,58 +156,52 @@ class Social_Media_Captioner:
 
 
     def _load_model(self):
-        try:
-            self.bnb_config = BitsAndBytesConfig(
-                load_in_4bit = True,
-                bnb_4bit_use_double_quant = True,
-                bnb_4bit_quant_type= "nf4",
-                bnb_4bit_compute_dtype=torch.bfloat16,
-                )
+        self.bnb_config = BitsAndBytesConfig(
+            load_in_4bit = True,
+            bnb_4bit_use_double_quant = True,
+            bnb_4bit_quant_type= "nf4",
+            bnb_4bit_compute_dtype=torch.bfloat16,
+            )
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.MODEL_NAME,
+            device_map = "auto",
+            trust_remote_code = True,
+            quantization_config = self.bnb_config
+            )
+
+        # Defining the tokenizers
+        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.MODEL_NAME,
-                device_map = "auto",
-                trust_remote_code = True,
-                quantization_config = self.bnb_config
+                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.MODEL_NAME)
+            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")
-
-        except Exception as e:
-            print(e)
-            self.model_loaded = False
-
+        self.model_loaded = True
+        print("Model Loaded successfully")
 
     def inference(self, input_text: str, use_cached=True, cache_generation=True) -> str | None:
         if not self.model_loaded: