Update handler.py

handler.py

@@ -1,5 +1,5 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-import torch # เพิ่มการ import torch
+import torch
 import os
 
 class EndpointHandler:
@@ -16,7 +16,8 @@ class EndpointHandler:
  model_name, 
  device_map="auto", 
  torch_dtype=torch.float16, 
- offload_folder=offload_dir
+ offload_folder=offload_dir,
+ offload_state_dict=True # เพิ่มพารามิเตอร์นี้เพื่อจัดการการ offload ให้ดียิ่งขึ้น
  )
  self.tokenizer = AutoTokenizer.from_pretrained(model_name)