--- library_name: transformers license: apache-2.0 --- # 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](https://huggingface.co/ThaiLLM/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](https://github.com/axolotl-ai-cloud/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 | ```bash 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 ```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 ```python 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: ลักษณะอุจจาระเป็นอย่างไร ถ่ายเป็นน้ำ เนื้อปนน้ำ มีมูกเลือด หรือเลือดสด?""" ``` ```python 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) ```