docs: adjust demo

README.md

@@ -70,11 +70,27 @@ tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=False, rev
 model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=False, dtype=torch.bfloat16, revision=revision).to(device)
 model.eval()
 
-
-prompt = "Die Altstadt von München "
-generator = pipeline('text-generation', model=model, tokenizer=tokenizer, device=device, max_new_tokens=max_new_tokens)
-outputs = generator(prompt)
-print(outputs)
+conversation = [
+    {"role": "user", "content": "What is the capital of France?"},
+]
+
+inputs = tokenizer.apply_chat_template(
+    conversation,
+    add_generation_prompt=True,
+    tokenize=True,
+    return_dict=True,
+    return_tensors="pt"
+).to(device)
+
+with torch.no_grad():
+    outputs = model.generate(
+        **inputs,  # Unpack the dictionary
+        max_new_tokens=args.max_new_tokens,
+    )
+
+# Decode only the generated tokens (excluding the input prompt)
+generated_tokens = outputs[0, inputs["input_ids"].shape[1]:]
+print(tokenizer.decode(generated_tokens, skip_special_tokens=True))
 ```
 
 ## License