Update model.py

model.py

@@ -1,5 +1,6 @@
 import os
 import torch
+import json
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 MODEL_NAME = "deepseek-ai/deepseek-llm-7b-chat"
@@ -16,17 +17,18 @@ model = AutoModelForCausalLM.from_pretrained(
 
 def classify_esi(symptoms):
     """
-    ใช้โมเดล AI ประเมินระดับ ESI จากข้อความอาการของผู้ป่วย
+    ใช้โมเดล AI ประเมินระดับ ESI จากข้อความอาการของผู้ป่วย โดยให้ตอบเป็น JSON
     """
     prompt = (
-        "คุณเป็นแพทย์เวชศาสตร์ฉุกเฉินที่ต้องประเมิน ESI 1-5 ตามอาการของผู้ป่วย:\n"
+        "คุณเป็นแพทย์ฉุกเฉินที่ต้องประเมินระดับ ESI 1-5 ตามอาการของผู้ป่วย:\n"
         "- ESI 1: ต้องได้รับการช่วยชีวิตทันที (หมดสติ, หัวใจหยุดเต้น, หายใจไม่ออก)\n"
         "- ESI 2: มีภาวะเสี่ยงสูง (แน่นหน้าอก, อาเจียนเป็นเลือด, ซึมมากผิดปกติ)\n"
         "- ESI 3: ต้องใช้ทรัพยากรทางการแพทย์หลายอย่าง (ตรวจเลือด, X-ray, ให้น้ำเกลือ)\n"
         "- ESI 4: ต้องใช้ทรัพยากรทางการแพทย์เพียงอย่างเดียว (เย็บแผล, ทำแผล, ฉีดยา)\n"
-        "- ESI 5: ไม่ต้องใช้ทรัพยากรทางการแพทย์ (แผลถลอก, น้ำมูกไหล, ไอเล็กน้อย)\n"
-        f"\nอาการของผู้ป่วย: {symptoms}\n"
-        "กรุณาตอบกลับเป็นตัวเลข 1, 2, 3, 4 หรือ 5 เท่านั้น (ห้ามตอบเป็นข้อความ)"
+        "- 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")
@@ -34,15 +36,18 @@ def classify_esi(symptoms):
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=10,  # ✅ เพิ่ม token ให้โมเดลมีพื้นที่ตอบกลับ
+            max_new_tokens=20,  # ✅ ให้พื้นที่โมเดลตอบ JSON
             do_sample=False  # ✅ ปิดการสุ่ม ลดโอกาสตอบผิด
         )
 
     response = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
 
-    # ✅ ตรวจจับตัวเลขที่เป็นไปได้
-    for esi in ["1", "2", "3", "4", "5"]:
-        if esi in response:
-            return int(esi)  # ✅ ส่งค่าที่ถูกต้องกลับไป
+    # ✅ พยายามดึง 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 "ไม่สามารถประเมินได้"