model.py

import os
import torch
import json
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_NAME = "deepseek-ai/deepseek-llm-7b-chat"
MODEL_DIR = os.getenv("HF_HOME", "/tmp")

# ✅ โหลดโมเดลและ Tokenizer ครั้งเดียว
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, cache_dir=MODEL_DIR)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16,
    device_map="auto",
    cache_dir=MODEL_DIR
)

def classify_esi(symptoms):
    """
    ใช้โมเดล AI ประเมินระดับ ESI จากข้อความอาการของผู้ป่วย โดยให้ตอบเป็น JSON
    """
    prompt = (
        "คุณเป็นแพทย์ฉุกเฉินที่ต้องประเมินระดับ ESI 1-5 ตามอาการของผู้ป่วย:\n"
        "- ESI 1: ต้องได้รับการช่วยชีวิตทันที (หมดสติ, หัวใจหยุดเต้น, หายใจไม่ออก)\n"
        "- ESI 2: มีภาวะเสี่ยงสูง (แน่นหน้าอก, อาเจียนเป็นเลือด, ซึมมากผิดปกติ)\n"
        "- ESI 3: ต้องใช้ทรัพยากรทางการแพทย์หลายอย่าง (ตรวจเลือด, X-ray, ให้น้ำเกลือ)\n"
        "- ESI 4: ต้องใช้ทรัพยากรทางการแพทย์เพียงอย่างเดียว (เย็บแผล, ทำแผล, ฉีดยา)\n"
        "- ESI 5: ไม่ต้องใช้ทรัพยากรทางการแพทย์ (แผลถลอก, น้ำมูกไหล, ไอเล็กน้อย)\n\n"
        f"อาการของผู้ป่วย: {symptoms}\n\n"
        "ตอบกลับเป็น JSON เท่านั้น เช่น:\n"
        "{ \"esi\": 3 }"
    )

    inputs = tokenizer(prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=20,  # ✅ ให้พื้นที่โมเดลตอบ JSON
            do_sample=False  # ✅ ปิดการสุ่ม ลดโอกาสตอบผิด
        )

    response = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()

    # ✅ พยายามดึง JSON ออกมา
    try:
        esi_data = json.loads(response)
        if "esi" in esi_data and isinstance(esi_data["esi"], int):
            return esi_data["esi"]
    except json.JSONDecodeError:
        pass

    return "ไม่สามารถประเมินได้"