Update handler.py

handler.py

@@ -2,7 +2,7 @@ import os
 from typing import Dict, List, Any
 from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
 import torch
-from unsloth import FastLanguageModel  # Import Unsloth
+from unsloth import FastLanguageModel
 from dotenv import load_dotenv
 
 load_dotenv()
@@ -16,57 +16,55 @@ class EndpointHandler:
         max_seq_length = int(os.getenv("MAX_SEQ_LENGTH", 2048))
         max_new_tokens = int(os.getenv("MAX_NEW_TOKENS", 512))
         self.hf_token = os.getenv("HUGGINGFACE_TOKEN")
-        self.model_dir = os.getenv("MODEL_DIR", ".") # Ensure this is set correctly
+        self.model_dir = os.getenv("MODEL_DIR", ".")
 
-        print(f"MODEL_DIR: {self.model_dir}")  # Debug print
-        print(f"Files in model directory: {os.listdir(self.model_dir)}") # VERY IMPORTANT
+        print(f"MODEL_DIR: {self.model_dir}")
+        print(f"Files in model directory: {os.listdir(self.model_dir)}")
 
         # --- 1. Load Config ---
-        # Load the configuration first to determine the base model type
-        self.config = AutoConfig.from_pretrained(self.model_dir, token=self.hf_token)
+        # Load the configuration first, WITH trust_remote_code=True
+        self.config = AutoConfig.from_pretrained(self.model_dir, token=self.hf_token, trust_remote_code=True)
 
         # --- 2. Load Tokenizer ---
-        #  Load the tokenizer.  Handle tokenizer loading errors.
         try:
-            self.tokenizer = AutoTokenizer.from_pretrained(self.model_dir, token=self.hf_token)
+            self.tokenizer = AutoTokenizer.from_pretrained(self.model_dir, token=self.hf_token, trust_remote_code=True)
         except Exception as e:
             print(f"Error loading tokenizer: {e}")
             raise
 
         # --- 3. Load Model ---
-        # Load the base model *and* apply the LoRA adapter.  Handle model loading.
         try:
-            # Load base model (using config to determine correct class)
+            # Load base model, WITH trust_remote_code=True
             self.model = AutoModelForCausalLM.from_pretrained(
-                self.config.base_model_name_or_path,  # Use base model from config!
+                self.config.base_model_name_or_path,
                 config=self.config,
-                torch_dtype=torch.bfloat16,  # Use bfloat16 if available
+                torch_dtype=torch.bfloat16,
                 token=self.hf_token,
-                device_map="auto",  # Let transformers handle device placement
+                device_map="auto",
+                trust_remote_code=True, #CRUCIAL
             )
 
             # Load and apply the LoRA adapter
-            self.model = FastLanguageModel.get_peft_model(self.model, self.model_dir) #model_dir contains lora weights
-            FastLanguageModel.for_inference(self.model) #Unsloth speed up
+            self.model = FastLanguageModel.get_peft_model(self.model, self.model_dir)
+            FastLanguageModel.for_inference(self.model)
 
         except Exception as e:
             print(f"Error loading model: {e}")
             raise
 
-                # Define the prompt style
-        self.prompt_style = """Below is an instruction that describes a task, paired with an input that provides further context. 
-Write a response that appropriately completes the request. 
+        # Define the prompt style
+        self.prompt_style = """Below is an instruction that describes a task, paired with an input that provides further context.
+Write a response that appropriately completes the request.
 Before answering, think carefully about the question and create a step-by-step chain of thoughts to ensure a logical and accurate response.
 ### Instruction:
-You are BuildwellAI, an AI assistant specialized in UK building regulations and construction standards. You provide accurate, helpful information about building codes, construction best practices, and regulatory compliance in the UK. 
-Always be professional and precise in your responses. 
+You are BuildwellAI, an AI assistant specialized in UK building regulations and construction standards. You provide accurate, helpful information about building codes, construction best practices, and regulatory compliance in the UK.
+Always be professional and precise in your responses.
+
 ### Question:
 {}
 ### Response:
 <think>{}"""
 
-
-
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
         Processes the input and generates a response.
@@ -84,8 +82,7 @@ Always be professional and precise in your responses.
         if torch.cuda.is_available():
             input_tokens = input_tokens.to("cuda")
 
-
-        with torch.no_grad():  # Ensure no gradient calculation
+        with torch.no_grad():
             output_tokens = self.model.generate(
                 input_ids=input_tokens.input_ids,
                 attention_mask=input_tokens.attention_mask,