handler.py

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

class EndpointHandler:
    def __init__(self, path=""):
        # ระบุชื่อโมเดลใน Hugging Face Hub
        model_name = "niruemon/llm-swp"

        # โหลดโมเดลและ tokenizer
        self.model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16)
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)

        # สร้าง pipeline สำหรับการสร้างข้อความ
        self.generator = pipeline("text-generation", model=self.model, tokenizer=self.tokenizer, device_map="auto")

    def __call__(self, data):
        # รับข้อความ input จากผู้ใช้
        input_text = data.get("inputs", "")
        if not input_text:
            return {"error": "No input text provided."}

        # สร้างข้อความโดยใช้โมเดล
        try:
            result = self.generator(input_text, max_length=150, num_return_sequences=1)
            generated_text = result[0]["generated_text"]
            return {"generated_text": generated_text}
        except Exception as e:
            return {"error": str(e)}