File size: 1,842 Bytes
0af61f8 3251ecc d81f76b 0af61f8 36f9031 0af61f8 d81f76b 36f9031 d81f76b 3251ecc d81f76b 36f9031 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
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,
offload_folder=offload_dir,
offload_state_dict=True # เพิ่มพารามิเตอร์นี้เพื่อจัดการการ offload ให้ดียิ่งขึ้น
)
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)}
|