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 peft import PeftModel  # Import PeftModel
+from peft import PeftModel
 from dotenv import load_dotenv
 
 load_dotenv()
@@ -12,64 +12,60 @@ class EndpointHandler:
         """
         Initializes the model and tokenizer.
         """
-        # Key settings (from environment variables, with defaults)
         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", ".")
-        self.base_model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B" # The base model ID
+        self.model_dir = os.getenv("MODEL_DIR", ".")  # Should be "." for root
+        self.base_model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"  # Base model!
 
         print(f"MODEL_DIR: {self.model_dir}")
         print(f"Files in model directory: {os.listdir(self.model_dir)}")
 
-        # --- 1. Load Config ---
-        self.config = AutoConfig.from_pretrained(self.base_model_name, token=self.hf_token, trust_remote_code=True)
+        # Load Config (with trust_remote_code)
+        self.config = AutoConfig.from_pretrained(
+            self.base_model_name, token=self.hf_token, trust_remote_code=True
+        )
 
-        # --- 2. Load Tokenizer ---
+        # Load Tokenizer (with trust_remote_code)
         try:
-            self.tokenizer = AutoTokenizer.from_pretrained(self.base_model_name, token=self.hf_token, trust_remote_code=True)
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                self.base_model_name, token=self.hf_token, trust_remote_code=True
+            )
         except Exception as e:
             print(f"Error loading tokenizer: {e}")
             raise
 
-        # --- 3. Load Model with PeftModel ---
+        # Load Model and LoRA Adapter (with trust_remote_code)
         try:
-            # Load base model
-             base_model = AutoModelForCausalLM.from_pretrained(
+            base_model = AutoModelForCausalLM.from_pretrained(
                 self.base_model_name,
                 config=self.config,
-                torch_dtype=torch.bfloat16,
+                torch_dtype=torch.bfloat16,  # Use bfloat16
                 token=self.hf_token,
                 device_map="auto",
-                trust_remote_code=True,
+                trust_remote_code=True,  # Important for Qwen2
             )
-
-            # Load the LoRA model using PeftModel
             self.model = PeftModel.from_pretrained(base_model, self.model_dir)
-            # No need for FastLanguageModel.for_inference() here, PeftModel handles it
 
         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. 
+        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.
-        """
         inputs = data.pop("inputs", None)
         if inputs is None:
             return [{"error": "No input provided. 'inputs' key missing."}]
@@ -77,14 +73,11 @@ Always be professional and precise in your responses.
             return [{"error": "Invalid input type. 'inputs' must be a string."}]
 
         input_text = self.prompt_style.format(inputs, "")
-
-        # Tokenize and move to CUDA (if available)
         input_tokens = self.tokenizer([input_text], return_tensors="pt")
         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,