Update model.py

model.py

@@ -1,12 +1,12 @@
 import os
 import torch
-import json
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
+# ใช้ Hugging Face Cache Directory
 MODEL_NAME = "deepseek-ai/deepseek-llm-7b-chat"
 MODEL_DIR = os.getenv("HF_HOME", "/tmp")
 
-# ✅ โหลดโมเดลและ Tokenizer ครั้งเดียว
+# ✅ โหลด Tokenizer และ Model ครั้งเดียวเท่านั้น
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, cache_dir=MODEL_DIR)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
@@ -17,18 +17,12 @@ model = AutoModelForCausalLM.from_pretrained(
 
 def classify_esi(symptoms):
     """
-    ใช้โมเดล AI ประเมินระดับ ESI จากข้อความอาการของผู้ป่วย โดยให้ตอบเป็น JSON
+    ใช้โมเดล AI ประเมินระดับ ESI จากข้อความอาการของผู้ป่วย
     """
     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 }"
+        "คุณเป็นแพทย์เวชศาสตร์ฉุกเฉินที่ต้องประเมิน ESI 1-5 ตามอาการที่ได้รับ..."
+        f"\nอาการของผู้ป่วย: {symptoms}\n"
+        "กรุณาประเมิน ESI เป็นตัวเลข 1-5 เท่านั้น ถ้าข้อมูลไม่พอให้ตอบว่า 'ไม่สามารถประเมินได้'"
     )
 
     inputs = tokenizer(prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
@@ -36,18 +30,14 @@ def classify_esi(symptoms):
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=20,  # ✅ ให้พื้นที่โมเดลตอบ JSON
-            do_sample=False  # ✅ ปิดการสุ่ม ลดโอกาสตอบผิด
+            max_new_tokens=5,  # ✅ จำกัดคำตอบให้สั้น
+            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
+    for esi in ["1", "2", "3", "4", "5"]:
+        if esi in response:
+            return int(esi)
 
     return "ไม่สามารถประเมินได้"