Update usage code example

README.md

@@ -23,7 +23,7 @@ pip install transformers==4.49.0 torch==2.6.0
 ```
 
 ```python
-from transformers import pipeline, TextStreamer, AutoModelForCausalLM, AutoTokenizer
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, TextStreamer
 import torch
 
 model_path = "Felladrin/Qwen2-96M"
@@ -32,15 +32,19 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 tokenizer = AutoTokenizer.from_pretrained(model_path)
 model = AutoModelForCausalLM.from_pretrained(model_path).to(device)
 streamer = TextStreamer(tokenizer)
-generate = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device)
-inputs = tokenizer(prompt, return_tensors="pt").to(device)
-model.generate(
-    inputs.input_ids,
-    attention_mask=inputs.attention_mask,
+generate = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    device=device,
     streamer=streamer,
-    max_length=tokenizer.model_max_length,
+)
+generate(
+    prompt,
     eos_token_id=tokenizer.eos_token_id,
     pad_token_id=tokenizer.pad_token_id,
+    max_length=tokenizer.model_max_length,
+    truncation=True,
     do_sample=True,
     repetition_penalty=1.05,
 )