Update handler.py

handler.py

@@ -1,28 +1,32 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from unsloth import FastLanguageModel
 import torch
 import os
 
 class EndpointHandler:
     def __init__(self, path=""):
-        # ระบุชื่อโมเดลใน Hugging Face Hub
-        model_name = "niruemon/llm-swp"
+        # ระบุชื่อโมเดลที่คุณต้องการใช้งาน
+        model_name = "defog/llama-3-sqlcoder-8b"
+
+        # Configuration settings
+        max_seq_length = 2048
+        dtype = None  # Keep as None, you can change later if needed
+        load_in_4bit = True
 
         # กำหนดไดเรกทอรีสำหรับการ 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, 
-            offload_state_dict=True,  # เปิดใช้งานการ offload state dict
-            offload_dir=offload_dir  # ระบุ `offload_dir` โดยตรงเพื่อจัดการการ offload ให้ถูกต้อง
+        self.model, self.tokenizer = FastLanguageModel.from_pretrained(
+            model_name=model_name,
+            max_seq_length=max_seq_length,
+            dtype=dtype,
+            load_in_4bit=load_in_4bit,
+            offload_folder=offload_dir  # ระบุโฟลเดอร์สำหรับการ offload
         )
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-        # สร้าง pipeline สำหรับการสร้างข้อความ
-        self.generator = pipeline("text-generation", model=self.model, tokenizer=self.tokenizer, device_map="auto", offload_folder=offload_dir)
+        # เตรียมโมเดลสำหรับการประมวลผลข้อความ
+        FastLanguageModel.for_inference(self.model)
 
     def __call__(self, data):
         # รับข้อความ input จากผู้ใช้
@@ -32,8 +36,9 @@ class EndpointHandler:
 
         # สร้างข้อความโดยใช้โมเดล
         try:
-            result = self.generator(input_text, max_length=150, num_return_sequences=1)
-            generated_text = result[0]["generated_text"]
+            inputs = self.tokenizer(input_text, return_tensors="pt").to(self.model.device)
+            outputs = self.model.generate(**inputs, max_new_tokens=150)
+            generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
             return {"generated_text": generated_text}
         except Exception as e:
             return {"error": str(e)}