Update README.md

README.md

@@ -124,35 +124,51 @@ This section provides a step-by-step guide to loading and using the model for ge
 pip install transformers torch peft
 ```
 
-## Import Libraries
+## Import Required Libraries
 
 ```python
-import os
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 ```
 
-## Load Model with Disk Offloading
+In this cell, we import the necessary libraries from the Hugging Face Transformers package.
+
+## Configure 4-bit Quantization
 
 ```python
-# Create a directory for offloading model weights
-os.makedirs("offload_dir", exist_ok=True)
+# Configure quantization settings
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype="float16",
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4"
+)
+```
+
+This cell configures 4-bit quantization to reduce the memory footprint of the model.
+
+## Load Model and Tokenizer
 
-base_model_name = "amornpan/V3_qwen2.5-32b-med-thai-optimized"
+```python
+# Load the model directly
+model_name = "amornpan/V3_qwen2.5-32b-med-thai-optimized"
 model = AutoModelForCausalLM.from_pretrained(
-    base_model_name,
+    model_name,
+    quantization_config=bnb_config,
     device_map="auto",
-    trust_remote_code=True,
-    offload_folder="offload_dir"  # Specify the offload directory
+    trust_remote_code=True
 )
-tokenizer = AutoTokenizer.from_pretrained(base_model_name, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 ```
 
-## Define Test Function
+In this cell, we load the pre-trained Qwen 2.5 32B Thai medical model.
+
+## Create Testing Function
 
 ```python
+# Function to test the model
 def test_model(prompt, max_new_tokens=256):
     system_prompt = "You are a question answering assistant. Answer the question as truthful and helpful as possible. คุณคือผู้ช่วยตอบคำถาม จงตอบคำถามอย่างถูกต้องและมีประโยชน์ที่สุด"
-    full_prompt = f"<s><|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
+    full_prompt = f"system\n{system_prompt}\nuser\n{prompt}\nassistant\n"
     
     inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
     
@@ -166,16 +182,18 @@ def test_model(prompt, max_new_tokens=256):
     )
     
     response = tokenizer.decode(generated_ids[0], skip_special_tokens=False)
-    
-    assistant_response = response.split("<|im_start|>assistant\n")[-1].split("<|im_end|>")[0]
+    # Extract only the assistant's response
+    assistant_response = response.split("assistant\n")[-1].split("")[0]
     
     return assistant_response
 ```
 
-## Test the Model
+This function handles the generation of responses.
+
+## Test with Example Question
 
 ```python
-# ทดสอบโมเดล
+# Test with a single example
 example_question = "อาการของโรคเบาหวานมีอะไรบ้าง"
 print(f"\nคำถาม: {example_question}")
 response = test_model(example_question)