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)}