Update README.md

README.md

@@ -24,21 +24,30 @@ pip install transformers accelerate peft
 ```
 
 Load the model.
-```{python}
-from transformers, import AutoTokenizer, AutoModelForCausalLM
+```python
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel, PeftConfig
 
 repo_id = "stefan-m-lenz/Qwen-2.5-7B-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")
+model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, 
+                                             device_map="auto")
 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“? Antworte nur kurz mit dem ICD-10 Code."""
+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,
@@ -48,8 +57,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)