Update README.md

README.md

@@ -30,20 +30,37 @@ pip install transformers accelerate peft
 
 Load the model.
 ```python
-from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 from peft import PeftModel, PeftConfig
 
 repo_id = "stefan-m-lenz/Mixtral-8x7B-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_4bit=True,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.bfloat16
+        )
+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“?"
+test_input = """Welche ICD-10-Kodierung wird für die Tumordiagnose "Bronchialkarzinom, Hauptbronchus" verwendet? Antworte nur mit dem ICD-10 Code."""
+
+input_str = tokenizer.apply_chat_template(
+    [{"role": "user", "content": test_input}],
+    tokenize=False,
+    add_generation_prompt=True,
+    enable_thinking=False
+)
 
 # Generate response
-inputs = tokenizer(test_input, return_tensors="pt").to("cuda")
+inputs = tokenizer(input_str, return_tensors="pt").to("cuda")
 outputs = model.generate(
     **inputs,
     max_new_tokens=7,
@@ -53,8 +70,8 @@ outputs = model.generate(
     top_p=None,
     top_k=None,
 )
-response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-response = response[len(test_input):].strip()
+generated_tokens = outputs[0, inputs["input_ids"].shape[1]:]
+response = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
 
 print("Test Input:", test_input)
 print("Model Response:", response)