Update README.md

README.md

@@ -30,14 +30,18 @@ pip install transformers accelerate peft
 
 Load the model.
 ```python
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 from peft import PeftModel, PeftConfig
 
 repo_id = "stefan-m-lenz/Mistral-Nemo-ICDOPS-QA-2024"
 config = PeftConfig.from_pretrained(repo_id, device_map="auto")
-model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, device_map="auto")
+quantization_config = BitsAndBytesConfig(load_in_8bit=True)
+model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, 
+                                             device_map="auto",
+                                             quantization_config=quantization_config)
 model = PeftModel.from_pretrained(model, repo_id, device_map="auto")
-tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path, device_map="auto")
+tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path,
+                                           device_map="auto")
 
 # Test input
 test_input = "Was ist der ICD-10-Code für die Tumordiagnose „Bronchialkarzinom, Hauptbronchus“?"