Update handler.py

handler.py

@@ -1,31 +1,28 @@
 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")
+class EndpointHandler:
+ def __init__(self, path=""):
+ # ระบุชื่อโมเดลใน Hugging Face Hub
+ model_name = "niruemon/llm-swp"
 
-# ประมวลผลคำขอจากผู้ใช้
-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)}
+ # โหลดโมเดลและ 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)}