ThaiLLM-27B-Prescreen

ThaiLLM-27B-Prescreen is a specialized model for initial screening of the disease the patient could possibly be suffering from, the department the patient should go to, and the severity of the condition. The model is a supervised fine-tuned model from ThaiLLM-27B.

Data

The model was trained on a chat style dataset, similar to that of the example in the usage below. The dataset comprises of 989 conversations that emulate a conversation betwen a patient and a doctor with the label being the disease, department, and severity that the doctor determines.

Performance

We evaluated our supervised-finetuned model against the base ThaiLLM-27B and Medgemma-27B-it, where we see that after SFT the ThaiLLM-27B-Prescreen rivals Medgemma-27B-it on our prescreen task.

Model	Department (F1)	Severity (F1)
ThaiLLM-27B (Zero-shot)	0.2320	0.2215
ThaiLLM-27B-Prescreen	0.3919	0.2186
Medgemma-27B-it (Zero-shot)	0.3853	0.3031

Training

The model was finetuned via axolotl with the following hyperparameters

Hyperparemeter	Value
Learning Rate	2e-4
LoRA Rank	16
LoRA Alpha	32
Sequence Length	2048
Epochs	3
Batch size	32

git clone https://github.com/axolotl-ai-cloud/axolotl.git

pip3 install -U packaging setuptools wheel ninja
pip3 install --no-build-isolation axolotl[flash-attn,deepspeed]

axolotl train prescreen.yaml

Prescreen.yaml

base_model: ThaiLLM/ThaiLLM-27B
plugins:
  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
strict: false

chat_template: gemma3
datasets:
  - path: prescreen.jsonl
    type: alpaca
output_dir: ./outputs/ThaiLLM-27B-Prescreen

sequence_len: 2048
sample_packing: true
ddp_find_unused_parameters: true

load_in_4bit: true
adapter: qlora
lora_r: 16
lora_alpha: 32
lora_target_modules:
  - q_proj
  - k_proj
  - v_proj
  - o_proj
  - down_proj
  - up_proj
lora_mlp_kernel: true
lora_qkv_kernel: true
lora_o_kernel: true

gradient_accumulation_steps: 8
micro_batch_size: 1
num_epochs: 3
optimizer: adamw_torch_4bit
lr_scheduler: cosine
learning_rate: 2e-4

bf16: auto
tf32: true

logging_steps: 1
flash_attention: true
warmup_ratio: 0.1
saves_per_epoch: 3
weight_decay: 0.01

Usage

The model expects a conversation between a patient and a doctor as the input

system_prompt = """# Instruction
You will receive a conversation between a doctor and a patient. The doctor asks diagnostic questions, and the patient provides responses about their symptoms, history, and condition.

## Your Task
Analyze the conversation and classify the case into three categories:
1. **Disease/Condition**: The specific medical condition the patient is suffering from
2. **Department**: The appropriate medical department for treatment
3. **Severity**: The urgency level of the case

## Classification Guidelines
### Disease/Condition
- Base your diagnosis on the symptom pattern, onset, duration, and clinical features
- Consider the most likely diagnosis given the presenting symptoms
### Department Selection
- Choose from: Dermatology, Emergency Medicine, Internal Medicine, Obstetrics-Gynecology, Ophthalmology, Orthopedics and Physical Therapy, Otolaryngology, Pediatrics, Psychiatry, Self-care / Observation, Surgery
### Severity Levels
- **Emergency**: Life-threatening conditions requiring immediate medical attention (severe pain >7/10, sudden onset of serious symptoms, signs of organ failure, severe bleeding, difficulty breathing, loss of consciousness)
- **Visit Hospital / Clinic**: Conditions requiring professional medical evaluation but not immediately life-threatening (moderate symptoms, persistent issues, need for diagnosis/treatment)
- **Observe at Home**: Minor conditions manageable with self-care and monitoring (mild symptoms, known conditions with clear management)

## Output Format
# Respond with ONLY a valid JSON object in this exact format:
# {\"disease\": \"exact disease name\", \"department\": \"exact department name\", \"severity\": \"Emergency|Visit Hospital / Clinic|Observe at Home\"}"""

prompt = """## Conversation
Doctor: เริ่มมีท้องเสียตั้งแต่เมื่อไร?
Patient: เริ่มมีอาการท้องเสียประมาณ 2 เดือนที่แล้วค่ะ
Doctor: มีอาการมานานเท่าไร?
Patient: ท้องเสียมาประมาณ 2 เดือนแล้วค่ะ
Doctor: ลักษณะอุจจาระเป็นอย่างไร ถ่ายเป็นน้ำ เนื้อปนน้ำ มีมูกเลือด หรือเลือดสด?"""

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

messages = [
    {"role": "system", "content": system_prompt},
    {"role": "user", "content": prompt}
]

tokenizer = AutoTokenizer.from_pretrained("ThaiLLM/ThaiLLM-27B-Prescreen", use_fast=True)
chat = tokenizer.apply_chat_template(prompt, add_generation_prompt=True, return_tensors="pt")
model = AutoModelForCausalLM.from_pretrained("ThaiLLM/ThaiLLM-27B-Prescreen", device_map="auto", dtype=torch.bfloat16)

output = model.generate(chat.input_ids, max_new_tokens=256)
response = tokenizer.batch_decode(output, skip_special_tokens=True)[0]
print(response)