Update handler.py

handler.py

@@ -1,13 +1,22 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-import torch
+import os
 
 class EndpointHandler:
     def __init__(self, path=""):
         # ระบุชื่อโมเดลใน Hugging Face Hub
         model_name = "niruemon/llm-swp"
 
+        # กำหนดไดเรกทอรีสำหรับการ offload โมเดล (สร้างขึ้นถ้ายังไม่มี)
+        offload_dir = "./offload"
+        os.makedirs(offload_dir, exist_ok=True)
+
         # โหลดโมเดลและ tokenizer
-        self.model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name, 
+            device_map="auto", 
+            torch_dtype=torch.float16, 
+            offload_folder=offload_dir
+        )
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
 
         # สร้าง pipeline สำหรับการสร้างข้อความ