Update handler.py

handler.py

@@ -1,8 +1,10 @@
 import os
+import time  # Import the time module
+import json
 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,22 +14,56 @@ 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", ".")  # Should be "." for root
-        self.base_model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"  # Base model!
+        self.model_dir = os.getenv("MODEL_DIR", ".")
+        self.base_model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"
 
         print(f"MODEL_DIR: {self.model_dir}")
-        print(f"Files in model directory: {os.listdir(self.model_dir)}")
+        print(f"Files in model directory (initial): {os.listdir(self.model_dir)}")
 
-        # Load Config (with trust_remote_code)
+        # --- Wait for adapter_config.json ---
+        config_path = os.path.join(self.model_dir, "adapter_config.json")
+        max_wait_time = 60  # Wait up to 60 seconds
+        wait_interval = 2  # Check every 2 seconds
+        start_time = time.time()
+
+        while not os.path.exists(config_path):
+            print(f"Waiting for adapter_config.json to appear...")
+            time.sleep(wait_interval)
+            if time.time() - start_time > max_wait_time:
+                raise FileNotFoundError(
+                    f"adapter_config.json not found after {max_wait_time} seconds."
+                )
+        print("adapter_config.json found!")
+
+        # --- Verify adapter_config.json contents ---
+        try:
+            with open(config_path, "r") as f:
+                adapter_config = json.load(f)
+                # Check for essential keys
+                if "base_model_name_or_path" not in adapter_config or \
+                   "task_type" not in adapter_config:
+                    raise ValueError("adapter_config.json is missing required keys.")
+                if adapter_config["base_model_name_or_path"] != self.base_model_name:
+                    raise ValueError("adapter_config.json base_model_name_or_path mismatch.")
+                if adapter_config["task_type"] != "CAUSAL_LM":
+                    raise ValueError("adapter_config.json task_type is incorrect.")
+                print("adapter_config.json contents verified.")
+
+        except (FileNotFoundError, json.JSONDecodeError, ValueError) as e:
+            raise Exception(f"Error verifying adapter_config.json: {e}")
+
+        print(f"Files in model directory (after wait): {os.listdir(self.model_dir)}")
+
+
+        # Load Config
         self.config = AutoConfig.from_pretrained(
             self.base_model_name, token=self.hf_token, trust_remote_code=True
         )
 
-        # Load Tokenizer (with trust_remote_code)
+        # Load Tokenizer
         try:
             self.tokenizer = AutoTokenizer.from_pretrained(
                 self.base_model_name, token=self.hf_token, trust_remote_code=True
@@ -36,15 +72,15 @@ class EndpointHandler:
             print(f"Error loading tokenizer: {e}")
             raise
 
-        # Load Model and LoRA Adapter (with trust_remote_code)
+        # Load Model
         try:
             base_model = AutoModelForCausalLM.from_pretrained(
                 self.base_model_name,
                 config=self.config,
-                torch_dtype=torch.bfloat16,  # Use bfloat16
+                torch_dtype=torch.bfloat16,
                 token=self.hf_token,
                 device_map="auto",
-                trust_remote_code=True,  # Important for Qwen2
+                trust_remote_code=True,
             )
             self.model = PeftModel.from_pretrained(base_model, self.model_dir)