Update model.py

model.py

@@ -2,11 +2,10 @@ import os
 import torch
 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 และ Model ครั้งเดียวเท่านั้น
+# ✅ โหลดโมเดลและ Tokenizer ครั้งเดียว
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, cache_dir=MODEL_DIR)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
@@ -20,9 +19,14 @@ def classify_esi(symptoms):
     ใช้โมเดล AI ประเมินระดับ ESI จากข้อความอาการของผู้ป่วย
     """
     prompt = (
-        "คุณเป็นแพทย์เวชศาสตร์ฉุกเฉินที่ต้องประเมิน ESI 1-5 ตามอาการที่ได้รับ..."
+        "คุณเป็นแพทย์เวชศาสตร์ฉุกเฉินที่ต้องประเมิน 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"
-        "กรุณาประเมิน ESI เป็นตัวเลข 1-5 เท่านั้น ถ้าข้อมูลไม่พอให้ตอบว่า 'ไม่สามารถประเมินได้'"
+        "กรุณาตอบกลับเป็นตัวเลข 1, 2, 3, 4 หรือ 5 เท่านั้น (ห้ามตอบเป็นข้อความ)"
     )
 
     inputs = tokenizer(prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
@@ -30,14 +34,15 @@ def classify_esi(symptoms):
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=5,  # ✅ จำกัดคำตอบให้สั้น
-            do_sample=False  # ✅ ปิดการสุ่มเพื่อลดเวลาประมวลผล
+            max_new_tokens=10,  # ✅ เพิ่ม token ให้โมเดลมีพื้นที่ตอบกลับ
+            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)
+            return int(esi)  # ✅ ส่งค่าที่ถูกต้องกลับไป
 
     return "ไม่สามารถประเมินได้"