Create handler.py

handler.py

@@ -0,0 +1,31 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import torch
+
+# โหลดโมเดลและ tokenizer
+def init():
+    global model, tokenizer, generator
+    model_name = "niruemon/llm-swp"  # ชื่อโมเดลใน Hugging Face Hub
+    
+    # โหลดโมเดลและ tokenizer
+    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    
+    # สร้าง pipeline สำหรับการสร้างข้อความ
+    generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device_map="auto")
+
+# ประมวลผลคำขอจากผู้ใช้
+def inference(inputs):
+    global generator
+    
+    # รับข้อความจาก inputs
+    prompt = inputs.get("text", "")
+    if not prompt:
+        return {"error": "No input text provided."}
+    
+    # สร้างข้อความโดยใช้โมเดล
+    try:
+        result = generator(prompt, 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)}